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)
Usage
This model converts an RGB image into a grayscale sketch using a PyTorch U-Net architecture.
Requirements
pip install torch torchvision pillow matplotlib
Load the model
The model checkpoint is available as:
best_sketch_unet.pt
First, define the model architecture:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
class DoubleConv(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, 3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.block(x)
class Down(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.block = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.block(x)
class Up(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.up = nn.ConvTranspose2d(
in_channels,
in_channels // 2,
kernel_size=2,
stride=2
)
self.conv = DoubleConv(
in_channels,
out_channels
)
def forward(self, x1, x2):
x1 = self.up(x1)
diff_y = x2.size(2) - x1.size(2)
diff_x = x2.size(3) - x1.size(3)
x1 = F.pad(
x1,
[
diff_x // 2,
diff_x - diff_x // 2,
diff_y // 2,
diff_y - diff_y // 2
]
)
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class SketchUNet(nn.Module):
def __init__(self):
super().__init__()
self.input = DoubleConv(3, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
self.bottleneck = Down(512, 1024)
self.up1 = Up(1024, 512)
self.up2 = Up(512, 256)
self.up3 = Up(256, 128)
self.up4 = Up(128, 64)
self.output = nn.Conv2d(64, 1, kernel_size=1)
def forward(self, x):
x1 = self.input(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.bottleneck(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
return torch.sigmoid(self.output(x))
Download and load the checkpoint
from huggingface_hub import hf_hub_download
checkpoint = hf_hub_download(
repo_id="Willie999/photo-to-sketch-unet",
filename="best_sketch_unet.pt"
)
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
model = SketchUNet().to(device)
model.load_state_dict(
torch.load(
checkpoint,
map_location=device
)
)
model.eval()
Generate a sketch
image = Image.open(
"your_image.jpg"
).convert("RGB")
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor()
])
input_tensor = transform(image)
input_tensor = input_tensor.unsqueeze(0).to(device)
with torch.no_grad():
sketch = model(input_tensor)
sketch = sketch[0, 0].cpu()
Display the result
plt.figure(figsize=(8, 8))
plt.imshow(
sketch,
cmap="gray"
)
plt.axis("off")
plt.show()
Save the generated sketch
output = (
sketch.clamp(0, 1) * 255
).byte()
output_image = Image.fromarray(
output.numpy()
)
output_image.save(
"generated_sketch.png"
)
Complete minimal example
Once the model architecture is defined, inference can be reduced to:
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms
import torch
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
checkpoint = hf_hub_download(
repo_id="Willie999/photo-to-sketch-unet",
filename="best_sketch_unet.pt"
)
model = SketchUNet().to(device)
model.load_state_dict(
torch.load(
checkpoint,
map_location=device
)
)
model.eval()
image = Image.open(
"your_image.jpg"
).convert("RGB")
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor()
])
x = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
sketch = model(x)
sketch = (
sketch[0, 0]
.cpu()
.clamp(0, 1)
.mul(255)
.byte()
)
Image.fromarray(
sketch.numpy()
).save("generated_sketch.png")
Input / Output
Input: RGB image
Output: Grayscale sketch
Input resolution: 256 × 256
Task: Image-to-image translation
Architecture: U-Net
Framework: PyTorch
The model is trained on paired photo/sketch examples from the rhfeiyang/photo-sketch-pair-500 dataset. It is intended as an experimental general image-to-sketch model and may perform differently depending on the subject, composition, and visual style of the input image.
