StawGAN: Structural-Aware Generative Adversarial Networks for Infrared Image Translation

The gist The proposed StawGAN introduces a novel structural-aware generative adversarial network to enhance the quality of translating night-time thermal infrared images into daytime color images by focusing on better-shaped and high-definition objects

Model Overview

The StawGAN is designed to translate night-time thermal infrared images (IR) into daytime color images (RGB) while specifically focusing on enhancing the quality of the target generation without merely colorizing it. The model is inspired by TarGAN [17] and incorporates a shared architecture with different flows for the image and the segmented target, interconnected by shared layers. The generator G takes as input the source domain image (s), the segmented image (rs), and a one-hot vector representing the target domain (t) to output both translated images and their segmentation.

Key Components and Losses

The network incorporates several novel components to achieve structural awareness, including a Shape Controller S that generates a binary mask representing the foreground area of the generated images. Additionally, a novel Contrast Network C is introduced to help the generator choose appropriate contrast, sharpness, and gamma factors. The training employs several loss functions to guide the generation process:

1. Adversarial loss for training both generators and discriminators.

2. A modality classification loss applied on the generator G to force the target modality.

3. A shape consistency loss Lshape-r, which tries to contain slight differences in structural shape by optimizing a Structural Dissimilarity Index Measure (DSSIM) loss encouraging spatial smoothness in the generated images.

4. The DSSIM loss is defined as Lssim = 1 − SSIM(xpaired − xt) where xpaired is the original image but in the same domain of xt, which is instead the translated sample.

5. A reconstruction loss Lrec-z aims at preserving the original characteristics of the source images, employing it with the DSSIM loss as suggested in [15].

6. A crossing loss Lcross is used to force the generator to pay attention to the target area when generating the translated image.

7. An L1 loss LL1 is applied, taking the fake image, with the applied factors, and the paired image serving in reducing production of more gray-like and more colorful images.

Experimental Setup and Results

The model was tested on the DroneVehicle [18] dataset, which contains 56, 878 samples divided between half RGB and half infrared image. Images are resized to 256 × 256 for training and validation, with preprocessing involving the removal of the darkest samples. Objective evaluation metrics include Frechet inception distance (FID) [24], Inception Score (IS) [25], structural similarity index (SSIM) [20], and PSNR. The results show that StawGAN improves state-of-the-art performance in most metrics compared to other models. Specifically, for the translation IR → RGB, StawGAN produces better samples with respect to other state-of-the-art models. Furthermore, the segmentation capability is also evaluated using Dice similarity coefficient (DSC) [26], Segmentation score (S-Score) [27], and mean absolute error (MAE). Table II confirms that StawGAN performs segmentation too, showing better performance with respect to TarGAN.

Conclusion

The proposed StawGAN introduces a neural approach that strengthens edges in the translation, making it useful for a variety of sub-tasks like object detection or classification. The method can be used with both paired or unpaired datasets, demonstrating its generalizability to different scenarios. The proposed model produces high-definition samples with realistic colorization and well-shaped and defined targets.

How it works

The encoder-decoder couples in the generator are composed of different building blocks: a convolutional one, with the usage of convolution, instance normalization, and Leaky ReLU, and an upsampling one involving the same operation as the previous one but with an upsampling operation done by interpolation at the beginning. The discriminator uses two identical models as separate entities to determine if an image is fake or real, with a conditional GAN for the segmented one to determine even if the domain is being correctly translated. The Shape Controller is built with convolutional blocks introduced before, while its upsampling is done via transposed convolutions. The Contrast Network consists of a convolutional layer, a max pooling, another convolutional layer, and finally three fully connected layers. This architecture allows the model to be target- and structural-aware by optimizing the Structural Dissimilarity Index Measure (DSSIM) loss.

Key Contributions

The main contribution is enforcing a better generation of targets using different flows in the same model for the image and the segmented target, with these flows being interconnected by a shared architecture. The method lies in the middle between paired and unpaired image translation models because it is generalizable to unpaired problems, but can also be used in paired ones taking advantage of ad hoc losses. The DSSIM loss is employed to improve the estimation of both the luminance and chrominance pixels of TIR images. This structural awareness produces better-shaped objects.

Comparison with State-of-the-Art

The comparison shows that StawGAN surpasses pix2pixHD in all other metrics, although its FID is higher than pix2pixHD's value. For the translation RGB → IR, both color and shape of the StawGAN image are more appealing with respect to pix2pixHD. TarGAN loses information on the targets since the image is darker, making it difficult to distinguish the shape and type of vehicle. The proposed StawGAN produces highdefinition samples with realistic colorization and well-shaped and defined targets. The comparison also shows that our model gains better performance with respect to TarGAN in segmentation.

References

[1] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio “Generative adversarial nets” in Advances in Neural Inf. Proc. Systems (NeurIPS), Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q.

1. The StawGAN model can produce better-shaped and high-definition objects in the target domain by enforcing structural awareness through a Structural Dissimilarity Index Measure (DSSIM) loss encouraging spatial smoothness in the generated images.

2. The improved AI system can perform accurate night-time thermal infrared image translation to daytime color images, specifically enhancing the quality of bus, car, etc. for subsequent tasks like object detection or classification, as it is designed to translate with attention to structural details.

3. The model's architecture incorporates a novel structure where different flows in the same model for the image and the segmented target are interconnected by shared layers, which strengthens edges in translation, making it useful for tasks requiring precise shape understanding.

4. The system can handle both paired and unpaired datasets because it is designed to be generalizable to unpaired problems, but can also be used in paired ones taking advantage of ad hoc losses.

5. The model utilizes a comprehensive loss function including adversarial loss, a modality classification loss to force the target modality, and structural losses like the DSSIM and shape consistency loss, ensuring both visual realism and accurate target structure preservation.

This paper addresses the problem of translating night-time thermal infrared images, which are the most adopted image modalities to analyze night-time scenes, to daytime color images (NTIT2DC), which provide better perceptions of objects. We introduce a novel model that focuses on enhancing the quality of the target generation without merely colorizing it. The proposed structural aware (StawGAN) enables the translation of better-shaped and high-definition objects in the target domain. We test our model on aerial images of the DroneVeichle dataset containing RGB-IR paired images. The proposed approach produces a more accurate translation with respect to other state-of-the-art image translation models. The source code is available at https://github.com/LuigiSigillo/StawGAN