Image colorization based on vision transformer

Hint-Based Image Colorization Based on Hierarchical Vision Transformer

Hint-based image colorization is an image-to-image translation task that aims at creating a full-color image from an input luminance image when a small set of color values for some pixels are given as hints. Though traditional deep-learning-based methods have been proposed in the literature, they are based on convolution neural networks (CNNs) that have strong spatial locality due to the convolution operations. This often causes non-trivial visual artifacts in the colorization results, such as false color and color bleeding artifacts. To overcome this limitation, this study proposes a vision transformer-based colorization network. The proposed hint-based colorization network has a hierarchical vision transformer architecture in the form of an encoder-decoder structure based on transformer blocks. As the proposed method uses the transformer blocks that can learn rich long-range dependency, it can achieve visually plausible colorization results, even with a small number of color hints. Through the verification experiments, the results reveal that the proposed transformer model outperforms the conventional CNN-based models. In addition, we qualitatively analyze the effect of the long-range dependency of the transformer model on hint-based image colorization. 

Paper link: https://www.mdpi.com/1424-8220/22/19/7419 

In this study, we proposed a vision transformer network with an encoder-decoder architecture for hint-based colorization. Through validation of the proposed model, we showed that the long-range dependency of the transformer can work effectively in hint-based colorization tasks. Even in regions where there were insufficient or no hints, the proposed model showed better results using color hints of similar objects in the image based on the long-range dependency. Furthermore, the experiments proved that the fewer hints, the greater the effect of the long-range dependency is. 

As shown in Figure 1, the proposed HCoTnet is divided into three main parts: first, a patch embedding (tokenization) module for using the input data (i.e., luminance image and color hint map) in the transformer; second, the Unet-like encoder and decoder modules consisting of transformer blocks [17]; finally, a projection module for outputting the result by restoring and projecting the embedded features onto the ab  dimension of the CIELAB color space. The luminance image and color hint map are the inputs of the network, and through the HCoT network, the colorization result is output with the ab  color channels.