Colorization of Grey Images by applying a Convolutional Autoencoder on the Jetson Nano

Tim Niklas Witte and Dennis Konkol

Contents

1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1 2 Convolutional Autoencoder . . . . . . . . . . . . . . . . . . . . . . . . . . . 1 3 Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1

3.1 Software . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1 3.2 Hardware . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2

4 Training . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2 4.1 Color Space Transformation . . . . . . . . . . . . . . . . . . . . . . . . . 2 4.2 Model . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2

5 Optimizing the model to run on the Jetson Nano . . . . . . . . . . . . . . 4 6 Evaluation: Compare with Colorful Image Colorization . . . . . . . . . . . 5 7 Conclusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7 References . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7

1 Introduction

Embedded GPU such as the Jetson Nano provide limited hardware resources compared with desktop/server GPUs. For example, the Jetson Nano has 128 CUDA cores and 4 GB video memory [1]. In comparison to the NVIDIA GeForce RTX 3070 Ti has 6144 CUDA cores and 8 GB video memory. As a consequence, inference done by huge artificial neural networks (ANN) i.e. over 21.000.000 parameter on the Jetson Nano becomes slow (ca. 1 forward pass per second). This paper presents a convolutional autoencoder for grey image colorization with 300.000 parameters optimized to run on embedded GPUs. In order to demonstrate the results during run time on the Jetson Nano, the live grey camera stream is colorized as shown in Figure 1.

Figure 1: OpenCV window on the Jetson Nano displaying the original, grey, colorized camera stream and corresponding loss between original and colorized image.

This paper is organized as follows: First of all, the concept of a convolutional autoencoder will be covered. Thereafter, the training procedure will be explained including the model architecture. Afterwards, our models performance is evaluated by comparing the colorized images generated by our models and by a state-of-the-art ANN for grey image colorization which has about 25.000.000 parameter. Finally, the final results are summed up.

2 Convolutional Autoencoder

3 Setup

3.1 Software

TensorFlow was installed following the official guide from NVIDIA. Furthermore, it is not recommened to install the current version of OpenCV via pip3 due to compatibility issues with the CSI camera. The CSI camera i.e. the gstream can only be accessed with a OpenCV version of lower than 3.3.1.

3.2 Hardware

As shown in Figure X, the CSI camera was plugged into the corresponding slot in the Jetson Nano. Furthermore, the HDMI display shows the OpenCV window as presented in Figure 1.

4 Training

Figure 2: Train and test loss during training.

4.1 Color Space Transformation

4.2 Model

As shown in Listing 1, our convolutional autoencoder has about 300.000 parameters. The models memory size is about 1.2 MB (300000 · 4 Byte). Encoder and decoder parts of the ANN are equality balanced due to having almost the same amount of parameters.

Listing 1: Parameter amount of our model (output of summary() call).

Model : ” autoencoder ”
Layer ( type )  Output Shape  Param #

=

encoder ( Encoder )  multiple  148155
decoder ( Decoder )  multiple  150145

=

Total params : 298 ,302
Trainable params : 297 ,210
Non−t r a i n a b l e params : 1 ,092

Figure 3 and 4 present the structure of the layers contained in the encoder and decoder. The encoder receives a 256x256 pixel grey image. Due to the grey color, there is only one color channel. Convolutions can be seen as a feature extractors. At first convolution in the encoder (see Conv2D 0 in Figure 3), there are 75 features extracted from this grey image. These extracted features are represented as channels (similar to color channel but not colors) called feature maps. Literally speaking, a feature map could be seen as a heatmap in which the pixel belonging to the corresponding feature has high values. Due to the stride size of 2, the size of these features maps is halved. A convolution operation is followed by a batch normalization layer and an activation layer (the drive is normalized before its goes into the activation function). In the encoder this ours four times. Each step the amount of filters increases.

Figure 3: Encoder layers.  Figure 4: Decoder layers.

The resulting embedding is passed into the decoder. Instead of convolutions reducing the feature map size, there are transpose convolutions which increase the feature map size by a factor of 2. Similar to the encoder, a transpose convolution is followed by a batch normalization layer and an activation layer. In the decoder this ours four times. Each step the amount of filters decreases. Except for the last transpose convolution which is a bottleneck layer: It decreases the amount of filters from 75 to 2 (a and b channel) and keeps the feature map size constant.

5 Optimizing the model to run on the Jetson Nano

Figure 5: Concept of a residual connection Residual connections also called skip connections in neural networks face the vanishing gradient problem (tiny weight adjustments) in the backpropagation algorithm. As shown in Figure 5, the output x of a layer is added two layers further as input to the third layer. The output x must be saved due to is used in a later time step. Therefore, residual connections need a lot of GPU memory causing a outsource of a part of other data needed for the model. For speed up the FPS, our model does not have residual connections.

As mentioned in the first section, the Jetson Nano has 128 CUDA cores. The amount of filters per layer does not exceed this number of cores. This limitation enables TensorFlow a simple scheduling of a filter map calculation to a specific core during output calculation of a layer In other words, there are no cores which must do a second filter map calculation after the first one while other cores are idling. The calculation of a previous layer must be finished before starting with the next layer. Furthermore, limit the amount of filer reduces the model size.

In Deep Learning overparameterization often occurs: The number of trainable parameters is much larger than the number of training examples. As a consequence, the model tends to overfit the data. The oppposite applies for our model. Literally speaking, our model is ”under-parameterizated” - Due to there are only 300.000 parameter on 1.3 million training images, our model is forced to generalize as strong as possible during training. To archive such generalization the model is training multiple epochs (iteration over entire training dataset). It is assumed that such generalization results into similar results compared with a model which as a multiple amount of parameters. In other words, the larger costs for training a small model compared with a larger model shall result into similar results but the latency to generate the result with the smaller model is lower. Besides, the non-existence of skip connections increases the chance of vanishing gradients during training. Although, this problem is compensated by multiple training epochs. To clarify, millions of tiny weight changes sum up into an effective weight adjustment.

6 Evaluation: Compare with Colorful Image Colorization

As demonstrated in Listing 2, the Colorful Image Colorization model from Richard Zhang et al. has about 25 million parameters. The model presented in this paper is about 80 times smaller. Its input shape 256x256x1 and the same as our model.

Listing 2: Parameter amount of the Colorful Image Colorization model (output of

        summary() call).

Model : ” C o l o r f u l I m a g e C o l o r i z a t i o n ”

Figure 6: Colorized images generated by the Colorful Image Colorization model from Richard Zhang et al. and by our model.

Layer ( type )  Output Shape  Param #

= [ …]

=

Total params : 24 ,793 ,081
Trainable params : 24 ,788 ,345
Non−t r a i n a b l e params : 4 ,736

Figure 6 shows grey images colorized by the Colorful Image Colorization model and by our model. Our model tends to colorize the images with a grey touch and the colors are not saturated compared with the Colorful Image Colorization model.

Our model does regression by predicting the ab values. The model output shape is 256x256x2 (see tanh 3 in Figure 4). In contrast to the model form Richard Zhang et al., classification is applied here: For pixel there is a probability distribution approximated which color it is. For demonstration purpose, there were 313 colors available. As a consequence, the model output shape is 256x256x313. The larger output shape compared to our model requires a larger (ca. 80 times) amount of parameters.

Figure 7: Loss based on colorized images by the Colorful Image Colorization model from Richard Zhang et al. and by our model.

Considering the loss as shown in Figure 7, our model outperforms the model form Richard Zhang et al. However, the Euclidean loss L_{2}(ŷ, y) for the prediction y and the target (also called ground truth) ŷ was applied:

L_{2}(ŷ, y) = \frac{1}{2} · ∑_{h,w} ||y_{h,w} − ŷ_{h,w}||^{2}

This loss function is ambiguous for the colorization problem. Consider the prediction y for a single pixel with a loss of d: There are two corresponding targets y ± d possible instead of a single one. Furthermore, consider a set of pixel, each of these pixel a corresponding color will be predicted. The optimal solution is the mean of all pixels within this set. In case of color prediction, this averaging causes a grey bias and desaturated colors.

7 Conclusion

Our model predicts the most possible color by applying regression. In contrast to the model proposed by Richard Zhang et al. which classifies the most possible color. Due to the one hot encoding applied for these color classification, there are over 80 times more parameter needed as required for our model considering the parameter balance between hidden layers and output layers. Comparing the colorized images generated by an ANN based on classification and by regression, regression based ANN tends to colorize images with a grey touch and unsaturated colors because of an ambiguous loss function for the colorization problem. However, the results are acceptable considering the difference in the amount of parameters between the two models. Furthermore, an classification based model can not be completely accelerated by a GPU because the last part of the model is a sampling process. This process is a argmax operation over a 313 possible colors (see model shape) which runs on the CPU. Note that, transferring data from GPU and CPU could be seen as a performance bottleneck.

Overall, our model archives about 10 FPS on the Jetson Nanon. Running the model of Richard Zhang et al. will results into less than 0.01 FPS.

References

[1] “Jetson nano developer kit.” https://developer.nvidia.com/embedded/jetson-nano-developer-kit. Accessed: 2022-03-24.