Photo-to-Sketch U-Net Model
This model is a custom PyTorch U-Net architecture (SketchUNet) trained to convert RGB photos into target outline sketches using the rhfeiyang/photo-sketch-pair-500 dataset.
Final Evaluation Metrics on Validation Set
- Combined Sketch Loss (L1 + Sobel Edge Loss): 0.14555
- Mean Absolute Error (L1 / MAE): 0.04962
- Mean Squared Error (MSE): 0.00653
- Structural Similarity Index (SSIM): 0.8804
- Peak Signal-to-Noise Ratio (PSNR): 22.68 dB
Hyperparameters
- Epochs: 50
- Batch Size: 8
- Optimizer: AdamW (Learning Rate: 2e-4, Weight Decay: 1e-4)
- Scheduler: CosineAnnealingLR (T_max: 50)
Intended Use Cases
This model can be integrated into various creative workflows:
- Digital Art & Graphic Design: Generate starting-point sketch outlines for painting.
- Mobile Photo Filters: Power interactive 'photo-to-sketch' features in camera apps.
- Prototyping & Storyboarding: Convert photographic scenes to monochrome mockups.
- Coloring Books: Automatically convert visual assets to line art.
Fast Inference (Hugging Face AutoModel)
You can perform high-speed inference locally using automated mixed-precision (autocast) and GPU acceleration:
import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoModel, AutoConfig
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = "Willie999/photo-to-sketch-unet"
# 1. Faster Initialization via Auto Classes
config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
model = AutoModel.from_pretrained(model_id, config=config, trust_remote_code=True).to(device)
model.eval()
# JIT trace or optimize compilation if doing batch inference
# model = torch.compile(model) # Uncomment on PyTorch 2.0+ for maximal speedup
# 2. Prepare the Input Image
image = Image.open("your_image.jpg").convert("RGB")
inference_transform = transforms.Compose([
transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.BICUBIC),
transforms.ToTensor(),
])
input_tensor = inference_transform(image).unsqueeze(0).to(device)
# 3. Fast Inference using torch.amp mixed precision autocasting
device_type = "cuda" if torch.cuda.is_available() else "cpu"
with torch.no_grad():
with torch.amp.autocast(device_type=device_type, dtype=torch.float16 if device_type == 'cuda' else torch.float32):
output = model(input_tensor)
# 4. Save the generated sketch
sketch_array = output[0, 0].cpu().numpy()
result_sketch = Image.fromarray((sketch_array * 255).astype('uint8'), mode='L')
result_sketch.save("output_sketch.png")