Color Constancy with Small Dataset via Pruning of CNN Filters

Abstract

Color constancy is an essential part of the Image Signal Processor (ISP) pipeline, which removes the color bias of the captured image generated by scene illumination. Recently, several supervised algorithms, including Convolutional Neural Networks (CNN)-based methods, have been proved to work correctly on this problem. However, they usually require a sufficient number of annotated data that reflects the complexity of realistic photometric effects and illumination in real scenes. It is time-consuming and costly to collect many raw images of various scenes with different lighting conditions and measure corresponding illumination values. The transfer learning technique, whose principal research focuses on collecting the knowledge gained when solving a problem and utilizing it to different but related problems, answers the need for large data. i.e., we transfer features from a CNN trained on a source task with a large scale dataset to the color constancy task. Most modern state-of-the-art CNN models are basically designed for the image classification task and are focused on training with deeper structures. One of the main disadvantages of deep convolutional neural networks is that they suffer from vanishing and exploding gradients. These deep structures also tend to overfit the data. Moreover, too many convolutional filters in these CNN models are not always beneficial to the networks, negatively influencing accuracy due to the many useless features. When it comes to the target task with small datasets, this redundancy is even worse. To reduce the dependence on a large scale labeled dataset and take advantage of standard and famous CNNs architectures, we proposed an approach to creating an efficient color constancy algorithm. Firstly, we utilized a structure channel pruning method named network slimming to thin our baseline model. It directly forces sparsity-induced regularization on the scaling factors in batch normalization layers, and less important channels are automatically distinguished during training and then pruned. Thereby, we iteratively pruned 75% channels of a specific Mobilenet version used as our model’s backbone, trained on a large-scale classification dataset. It means the backbone with the classification head is used to deal with our network pruning task. Experimental results show that our proposed model can efficiently adapt a new lighter backbone. However, the more parameters pruned away, the less the accuracy. Then the resulted compact model was transferred and trained on a small dataset doing color constancy. During training on the color constancy task, we applied the DSD technique. It regularized the network iteratively by learning connections’ importance during the initial dense solution and pruning the unimportant connections. Then, the pruned connections are recovered, and the whole network is retrained again. The proposed method reaches comparative performance with other state-of-the-art models, produces fewer MACs, and can significantly decrease computational costs.