From Sketch to Color with Pix2Pix: Our End-to-End Architecture, Training Strategy, and Deployment

By Muhammad Hassan and Muhammad Mashhood

Transforming a simple line sketch into a realistic, colored image is one of the most fascinating applications of deep learning in computer vision. Among the many image-to-image translation models, Pix2Pix remains one of the most practical and reliable baselines for supervised tasks of this kind.

In this project, we built a complete sketch-to-color generation pipeline using Pix2Pix, covering everything from model architecture and training strategy to deployment with an interactive web interface.

Our system includes:

• a U-Net based generator
• a conditional PatchGAN discriminator
• mixed precision training for faster and more efficient learning
• checkpointing and resume support
• a Gradio-powered deployment UI
• deployment through Hugging Face Spaces

This post walks through the entire system end to end—what we built, why each component matters, and how the complete pipeline works in practice.

Problem Setup

The task we address is a classic supervised image-to-image translation problem.

Given an input sketch, the model must generate its corresponding realistic colored version.

Objective

• Input: Sketch image
• Target: Colored image
• Goal: Learn a mapping from the sketch domain to the color image domain

In simpler terms, the model learns how to translate structural outlines and edges into meaningful textures, shades, and colors.

Because this is a supervised learning setup, every training sample must exist as a paired example:

one sketch image + its exact corresponding color target image

This pairing is essential because the model learns by comparing its generated output with the ground-truth target image.

Dataset Format

Our training pipeline supports two paired-data formats, making it flexible for different datasets.

1. Concatenated Paired Images

Both the sketch and target image are stored inside a single image file.

For example:

| sketch | color target |

This is one of the most common Pix2Pix dataset formats.

2. Separate Input-Target Pairs

The sketch and color images are stored as separate files with matching names.

Example:

input/
   img_001.png
target/
   img_001.png

This format is easier to manage for larger datasets and custom preprocessing pipelines.

Performance and Stability Improvements

We added several practical engineering improvements to make training faster and more stable.

Included optimizations

• automatic mixed precision (AMP) on CUDA
• optional multi-GPU DataParallel
• channels-last memory format
• config-driven dataloader settings
• resume-from-checkpoint support

These changes significantly improved throughput and reduced training interruptions.

Inference and Deployment

After training, we export the generator state_dict and use it for inference on uploaded sketches.

To make the model easily accessible, we wrapped inference in a Gradio interface and deployed it as an interactive application.

User workflow

• upload sketch
• select model/checkpoint
• generate colorized output

This created a fast transition from research experimentation to a shareable demo.

What We Learned

The biggest improvements in our system came from:

• U-Net skip connections for structural fidelity
• conditional discriminator for sketch-output alignment
• adversarial + L1 combined loss
• robust checkpointing
• fast Gradio deployment workflow

Each of these contributed significantly to both model quality and usability.

Final Thoughts

Pix2Pix continues to be a practical and powerful architecture for paired image translation when implemented carefully.

By combining:

• clean data pairing
• U-Net skip connections
• conditional GAN training
• robust checkpointing
• deployment-ready UI

we built a sketch-to-color system that is both trainable and production-friendly.

If you’re building a similar pipeline, this architecture remains an excellent baseline to start from and iterate upon.

Next Steps

Some directions we are currently exploring:

• logging adversarial and L1 losses separately over epochs
• side-by-side visual validation snapshots
• perceptual loss for richer texture quality
• checkpoint selector in UI for direct epoch comparison