Thermal infrared image colorization based on pre-trained stable diffusion and its shadow network with attention mechanism

Abstract Thermal infrared images are grayscale images obtained by measuring the intensity of infrared radiation. At night or in extreme weather, thermal imaging sensors have stronger ability of scene description in environmental monitoring. In recent years, The applications of measuring infrared radiation intensity and the limitations of thermal infrared images, such as low resolution, low contrast, and low signal-to-noise ratio, have prompted the demand for the colorization of thermal infrared images. In this work, we employ stable diffusion as the backbone to introduce an improved thermal infrared images colorization model, which includes a pre-trained network, a shadow network, and a fusion module to achieve better colorized results and reduce diversity of results from stable diffusion at the same time. First, we use stable diffusion model combined with its shadow network to generate realistically colorized images, the shadow network guides stable diffusion model while imposing constraints on its generated results, thus enhancing the overall generation performance. Second, we incorporate a fusion module between stable diffusion and its shadow network, the fusion module combines information from both stable diffusion and its shadow network, by computing Gram matrices and channel relationships. The fusion module strengthens important aspects of the generated color images while attenuating irrelevant parts. Furthermore, multi-scale convolution is used to enhance the performance of information extraction. Finally, we identify suitable loss functions to enhance the metrics. To validate the effectiveness of our model, we conduct evaluations on the KAIST and IRVI datasets. Experiments indicate that the approach we propose can generate more accurate and realistic color information compared with state-of-the-art performance. The SSIM increased by 0.09 on the KAIST dataset. On the IRVI dataset, SSIM increased by 0.05 and 0.01 for the two scenes, while PSNR increased by 0.39 and 0.68, respectively. By employing diffusion models to imbue infrared images with realistic colors, our approach substantially improves their readability and interpretability. This advancement not only facilitates a more intuitive human understanding of infrared data but also extends the application scope of diffusion models. More detailed information can be found at https://github.com/jinxinhuo/PTSD-for-Infrared-Image-Colorization