Generative / vision
Image Colorization: CNN vs U-Net vs CVAE
23.3 dB
PSNR for the CVAE (U-Net 21.5, CNN 20.0)
Colorization is ambiguous: a grey horse can be many colours. The comparison shows how skip connections sharpen output and how a latent variable lets the model express that ambiguity.
What I did
- Built a plain encoder-decoder CNN baseline, a U-Net with skip connections, and a conditional VAE trained on reconstruction plus a beta-weighted KL term.
- Results (test): CNN MSE 0.0099, SSIM 0.914; U-Net 0.0070, 0.938; CVAE 0.0046, 0.961. The CVAE also produces multiple plausible colourings.
- Wrote next steps I would take with more time: a VGG perceptual loss and a standard SSIM implementation.
Stack
PyTorchGradioCIFAR-10