Gray to Glory: Performance Analysis of Image Coloration Models

Approach

Data Preprocessing

• The images from the COCO dataset were converted into the Lab color space, where the L* channel (lightness) represents grayscale information, and the A* and B* channels encode color information.
• The L* channel was used as the input to the model, while the A* and B* channels served as the ground truth outputs for supervised learning.
• To ensure consistency, all images were resized to 256x256 pixels and normalized to a range between -1 and 1 for both input L channel and output AB channel.
• Random horizontal flipping was applied to the images before breakdown into L and AB channel to enhance model generalization and robustness.

Model Architectures

• CNN Encoder-Decoder: We started with a simple encoder-decoder CNN architecture [8] as our baseline. The encoder extracted latent grayscale features, while the decoder reconstructed the A* and B* color channels. This model served as a foundational benchmark for further comparisons.
  • Encoder: This module captures the spatial features by applying multiple convolution layers followed by downsampling like max pooling. This helps in reducing the dimensions and extracts important features.
  • Lantent Representation: The encoder outputs a feature representation which is compact in nature and gives the most vital spatial and semantic information
  • Decoder: The decoder takes the input from the encoder and performs upsampling with operations like transposed convolutions or interpolation. It thereby increases the spatial dimensions, performs convolutions to refine the features to output the finaloutput.
  • Output layer: We have a final convolution layer that outputs the desired output which is in the form of a 2-dimensional AB channel. This output is concatenated with the L channel to produce a LAB image which is then converted to an RGB image.


• U-Net with ResNet-18 Backbone: The U-Net architecture [6], with a ResNet-18 backbone in the encoder, was the second phase of our approach. Pre-trained weights from ImageNet were loaded into the encoder to leverage transfer learning, enhancing the model's ability to extract meaningful features. The pre-trained U-Net was fine-tuned on our dataset for 20 epochs to adapt to the specific task of image colorization. The U-Net model was implemented with the DynamicUnet class in fastai [9].
  • Encoder: Similar to the Encoder-Decoder module, the Encoder in UNet captures importatant spatial features, reducing the image size while extracting important features. Here we replace the traditional encoder with a ResNet-18 architecture. ResNet-18 consists of convolutional layers and skip connections (in the form of residual blocks) that help to learn deeper representations without having the vanishing gradient problem.
  • Connection: The deepest part of the UNet from the encoder gets connected to the decoder, which has captured the most important spatial features
  • Decoder: This model performs upsamling and recaptures the image resolution by concantenating them with corresponding feature maps from encoder.
  Skip Connections: The skips connections link the encoder layer to the corresponding decoder layers, which preserves fine-grained spatial features which get lost during the downsampling operation. Training: The UNet was trained using the Adam optimizer for 20 epochs with L1 loss as the criterion, since it - preserves spatial structure, it is robust to any outliers and has a better gradient flow.


• Conditional GAN: The pre-trained U-Net was integrated as the generator in a Conditional GAN (cGAN) framework. cGANs were chosen because they condition the output (colorized image) on the input (grayscale image), enabling the generator to produce outputs that are contextually consistent with the input image. Our implementation is inspired by the pix2pix implementation [4].
  • PatchGAN Discriminator: We used a PatchGAN discriminator, which evaluates the realism of small image patches rather than the entire image. This approach ensures local coherence, as the discriminator classifies each patch of the generated image as real or fake. PatchGAN was chosen because it is effective in preserving fine details and textures [5], which are critical for realistic image colorization.
  • Generator: The cGAN uses the UNet module as it's generator output. It creates images using the UNet module by effectively capturing spatial dimensions with skip layers. This enables accurate mappings from grayscale to colorized outputs.
  • Discriminator: This is a custom CNN module which differentiates between the real images or generated ones. This assesses the authenticity, which makes use of convolutions to differentiate between real or generated real and generated colorizations.
  • Training: We trained the cGAN using BCE Loss for 20 epochs. We used the adam optimizer again for training the discriminator with decay rate taken as 0.9.

Implementation Details

• Dataset: A subset of COCO dataset was used, consisting of 10,000 images. The dataset was split into 8,000 images for training and 2,000 for testing.
• Training Setup: The CNN encoder-decoder was trained for 100 epochs as a baseline. The U-Net architecture, pre-trained on ImageNet, was fine-tuned on our dataset for 20 epochs. The Conditional GAN was trained by integrating the this U-Net as a generator with the PatchGAN discriminator.
• Pre-trained Models: The ImageNet pre-trained weights for the ResNet-18 backbone was utilized for U-net generator pre-trainig to expedite the training process and to get better performance compared to a model trained from scratch.

Design Choices

Our GAN architecture is inspired by the pix2pix architecture [4].
• Why Conditional GAN: cGANs condition the output on the input, ensuring that the generated image aligns with the grayscale input's semantic content. Conditional GANs (cGANs) learn a conditional generative model [10]. This makes cGANs suitable for image-to-image translation tasks, where we condition on an input image and generate a corresponding output image.[4]
• Why PatchGAN Discriminator: PatchGAN discriminates at the level of image patches, which is effective for ensuring local coherence in generated images. This is particularly useful for image colorization, where local color consistency is critical.

Experiments and Results

Experimental Setup

• Datasets: A subset of the COCO dataset with 10,000 images. The training set consisted of 8,000 images, and the test set included 2,000 images.
• Metrics: L1 Error was used to evaluate pixel-wise accuracy during training. Mean Absolute Error (MAE) and Learned Perceptual Image Patch Similarity (LPIPS) score (using latent features from pretrained VGG) [7] are used to quantitatively compare the model performances.
• Hardware: The experiments were performed on Google Colab with the L4 GPU.

Results

• CNN Encoder-Decoder: Achieved baseline results with limited ability to capture complex features, showing noticeable artifacts in the colorized outputs. We got a validation loss of 0.0848 after training it for 100 epochs. We trained the module using the L1 loss criterion with learning rate set as 0.001
• U-Net with ResNet-18 Backbone: Improved performance compared to the baseline, with enhanced feature extraction and better contextual colorization. The UNet was trained for 20 epochs and acheieved a L1 loss of 0.07586. We set the learning rate to 0.0001
• Conditional GAN: Produced the most realistic and visually consistent outputs. The incorporation of adversarial loss with the PatchGAN discriminator significantly reduced artifacts and improved local color coherence. We achieved a discriminator loss of 0.59404 and discriminator loss of 11.37 after training for 20 epochs.

The MAE and LPIPS value on the test dataset for each model is listed in Table 1.

Table 1: Results

Model	MAE	LPIPS
CNN Encoder-Decoder	0.0848	0.148
U-Net (ResNet-18 backbone)	0.0846	0.141
cGAN with U-Net	0.0807	0.132

Qualitative Results

Parameter Analysis

• Epoch: We have tried to use the number of epochs for getting better accuracy but also making sure we do not overfit any model. Hence for training the UNet we used 20 epochs, after which we were seeing results getting converged with no frther reductions in losses. Similarly for cGAN, we kept the epoch numbers to 20 so that the generator does not overflow and discriminator does not overfit.
• Learning Rate: Since learning rate is a vital parameter in training neural networks since they control by how much the weights should be adjusted, we kept it as 0.01 so as to avoid any irregularities and unstable updates that would have destabilized the training.
• Beta1: Beta1 used in training the GAN is also a important parameter for the ADAM optimizer to take into account past weight updates for the current weights. We kept it as 0.5 to prevent any rapid updates, which allowed for stable training especially in earlier stages.
• Criterion: We used L1 loss as the criterion for training because it allows for pixelwise accuracy which helps in generating realistic images. We used the BCE loss for training the GAN which is a common method, the discriminator learns to differentiate between real and fake images and generator is given the taks to fool the discriminator into classfying the generated images as real ones.
• Batch Size: The batch size allows the neural network to update the weights in batches (controls the frequency as to when weights are updated). Since larger batch size required more memory and smaller ones lead to noiser gradients, we set the batch size as 32.

Trends

Conclusion

Future Improvements

• Data Augmentation: Expanding the dataset with more diverse images and applying advanced augmentation techniques could help the model generalize better to unseen images.
• Loss Functions: Experimenting with perceptual loss or combining MSE with adversarial loss might improve the model’s ability to generate outputs that are both quantitatively and qualitatively superior.
• Higher-Resolution Models: Incorporating super-resolution techniques or adapting the model for higher-resolution inputs could enhance the fine details in the colorized images.
• Diffusion Models: Diffusion models are state of the art image generation models which iteratively take noisy or partially incomplete images and refine them to produce a high-quality output.
  • Models like Denoising Diffusion Probabilistic Models (DDPMs) [11] can be used, which add noise to images in the feed forward process and learns to reverse this during the training period. DDPM and Guided Diffusion [12] are some framework that can be explored in future.
  • Another set of these models are Latent Diffusion Models (LDM) [13], which compress and compute the predictions in a compressed data space making them more efficient during computation. Stable Diffusion which are state of the art (SOTA) models nowadays, can be used to generate more realistic images.