Image is converted to lab space(2 channel) to reduce the regression by a channel. Model is trained on UNET architecture. Tried with various loss functions (MSE, SSIM, TVLOSS, Pretrained VGG feature loss). Weighted loss gave better performance.
image colorization
unet VGG feature loss Lab space image CNN regression Pytorch