Infrared-to-visible light image translation method based on comparative learning

Abstract

The invention provides an infrared-to-visible light image translation method based on comparative learning. An image data set and an image translation model are constructed, the model comprises an improved ViT feature extraction module, improved multilayer contrast loss and a generative adversarial network composed of a ResNet generator and a multi-scale discriminator, the network is trained by using a specific loss function, and then translation from an infrared image to a visible light image is completed by using the trained network. According to the method, the real visible light image with fine granularity can be generated from the infrared image, and the method can be applied to scenes such as night infrared-to-daytime visible light image translation, target detection and image fusion.

Description

Technical Field

[0001] The invention belongs to the technical field of image processing, and in particular relates to an infrared to visible light image translation method based on contrast learning. Background Art

[0002] Infrared images and visible light images are crucial imaging modalities in the field of computer vision. Visible light images can convey more visual information, such as color, texture, and morphology, which conforms to the perception characteristics of the human eye, but perform poorly at night or under extreme conditions. In contrast, infrared images use the principle of thermal imaging, have excellent imaging capabilities at night and in the dark, have good detection and recognition capabilities for objects that emit heat, and provide important spatial structure descriptions for visual signals. However, due to the lack of detailed semantic information in infrared images, it is difficult to directly apply them to practical tasks such as autonomous driving or monitoring.

[0003] Visible light images and infrared images taken at night are not satisfactory when used for tasks such as target detection, but visible light images generated from infrared images provide better results. Therefore, translating infrared images into visible light images can be used as an effective solution to improve the accuracy of target detection and assist unmanned driving projects, which has great research significance.

[0004] Infrared to visible light image translation can be regarded as a style transfer task, whose main goal is to convert infrared images into images with visible light style while ensuring the consistency of spatial structure. In this task, the key research difficulty lies in how to effectively constrain and maintain cross-domain invariance. CycleGAN uses two pairs of generators and discriminators, one pair is used to learn the mapping from the source domain to the target domain, and the other pair is used to learn the mapping from the target domain to the source domain. I2V-GAN introduces improved perceptual cycle loss and in-domain and out-of-domain similarity loss based on CycleGAN, which improves the detailed performance of generated visible light images, but the constraint effect and training time still need to be improved. There are still problems with the existing model. First, the need for two sets of generators and discriminators will lead to more computational costs, and second, its domain-invariant constraint ability is not strong enough. Summary of the invention

[0005] In order to overcome the shortcomings of the prior art, the present invention provides an infrared to visible light image translation method based on contrastive learning. An image data set and an image translation model are constructed. The model includes an improved ViT feature extraction module, an improved multi-layer contrast loss, and a generative adversarial network composed of a ResNet generator and a multi-scale discriminator. The network is trained using a specific loss function, and the trained network is used to complete the infrared image to visible light image translation. The present invention can generate a real, fine-grained visible light image from an infrared image, and can be applied to scenes such as nighttime infrared to daytime visible light image translation, target detection, and image fusion.

[0006] A method for infrared to visible light image translation based on contrastive learning, characterized by the following steps:

[0007] Step 1: construct an image dataset for training and testing an image translation model, wherein the image dataset includes infrared images and visible light images;

[0008] Step 2: construct an image translation model, the image translation model includes:

[0009] Improved ViT feature extraction module for image feature extraction and obtaining marker embedding;

[0010] Improved multi-layer contrast loss, used to calculate the mutual information of input and output images, using its maximum correlation as one of the loss functions;

[0011] The generative adversarial network composed of a ResNet generator and a multi-scale discriminator takes an infrared image as input and generates the corresponding visible light image.

[0012] Step 3, training the generative adversarial network, specifically comprising: sending the infrared image to the generator to generate a corresponding visible light image, and then sending the generated image to the discriminator to judge whether the generated image is true or false, and using adversarial loss and multi-layer contrast loss to train the generator and the discriminator;

[0013] Step 4: Input the infrared image into the trained generator to obtain a visible light image and complete the image translation.

[0014] Furthermore, in step 1, the InfraredCity-Lite and IRVI datasets dedicated to infrared-visible light image or video translation are used for model training and testing.

[0015] Furthermore, in step 2, an image translation model is constructed as follows:

[0016] Step 2-1, improved ViT feature extraction module: use the ViT model pre-trained on a large-scale dataset to perform feature extraction and obtain tag embedding. Specifically: given a two-dimensional image, the ViT model first divides the image into N patches, and then converts it into a vector sequence, which also includes category embedding and position embedding; then the image vector sequence is sent to the Transformer encoder. During the encoding process, the tag embedding of the n layers is first extracted, and then its category embedding is deleted as the final output of the n-layer feature, each layer of feature is a sequence with a length of N and a dimension of D; the encoder is composed of L encoding blocks, each encoding block is composed of a multi-head attention sublayer and an MLP sublayer, and L is set to 12; the D represents the sequence dimension, which is set to 512, and the n is an integer in the range of (1,6);

[0017] Step 2-2, improved multi-layer contrast loss: For the input image x∈X and its corresponding generated image Send them into the ViT model respectively to obtain the labeled embedding features of a total of I layers. The i-th layer labeled embedding features are expressed as ,i∈[1,I], let ViT model be V, which is represented as follows:

[0018]

[0019] Where X represents the input image dataset, Indicates the generation of image dataset;

[0020] for and , sample them at the same position to make their sequence length and dimension become K, and normalize the sampled sequence vector to a unit ball, and get exist Take a vector For query, Take the vector at the corresponding position is positive, the remaining vectors is negative, for z,z + and z - Apply the infoNCE loss function:

[0021]

[0022] Among them, l(z,z + ,z - ) represents z,z + and z - The infoNCE loss function value of the application, τ represents the temperature coefficient, and the value is τ=0.07; the K is set to 256;

[0023] Embed the i-th layer tag ,use Denotes the jth vector, j∈[1,K], which embeds the i-th layer tag ,use represents the jth vector, and the remaining vectors are represented as , the improved multi-layer contrast loss is calculated as follows:

[0024]

[0025] Among them, L MNCE Represents the multi-layer contrast loss value, Represents the mathematical expectation of a calculation;

[0026] Step 2-3, the generative adversarial network composed of the ResNet generator and the multi-scale discriminator: the ResNet generator is a residual generator, which inputs the input image and outputs the generated image; the multi-scale discriminator judges the authenticity of the generated image, as follows:

[0027] First, the real image y and the generated image Send them into the ViT model to get a layer of labeled embedding features Then reshape the label embedding into the same three-dimensional form set P as the image y , ,in Then, non-overlapping spatial averages are performed on adjacent vectors with receptive fields of different sizes, including After spatial averaging with a receptive field of 2×2, the set After the spatial average with a receptive field of 4×4, the set This process is continued until all spatial average sets of different scales are obtained; then, all the obtained sets are restored to the form of labeled embedding, including P y 2 , Restore to T y 2 , , P y 4 , Restore to T y 4 , , so that all the recovered sets are obtained, and they are concatenated to obtain and ; Finally, and They are sent to a three-layer learnable MLP classification network respectively, and the MLP outputs the result of whether it is true or false.

[0028] Furthermore, in step 3, the generative adversarial network is trained as follows: Given an infrared image dataset and visible light image datasets , the generator G learns the mapping relationship from X to Y, converts X into an image with the features of Y, but still maintains its original structure, and obtains the generated image The discriminator D distinguishes images And the real image sample y∈Y to output the result of judging whether the generated image is true or false;

[0029] The loss function is used to train the entire generative adversarial network, and the network is trained by minimizing the loss function of the discriminator and the loss function of the generator. The loss function of the discriminator L D for:

[0030]

[0031] Among them, L D represents the loss function value of the discriminator, G represents the generator, and D represents the multi-scale discriminator. MPL represents the MLP classification network, Indicates the real image feature set obtained by the intermediate calculation of the discriminator, Represents the generated image feature set obtained by the intermediate calculation of the discriminator;

[0032] The loss function of the generator is:

[0033]

[0034] Among them, L G represents the loss function value of the generator, is the adversarial loss, L MNCE is the multi-layer contrast loss, λ is the control L MNCE The weight of is in the range of (0,1).

[0035] The beneficial effects of the present invention are as follows: by constructing an image translation model, combining the feature extraction module of ViT with the classic image translation algorithm CUT, and designing an improved multi-layer contrast loss function, the spatial consistency in image translation can be better maintained, the scene structure information can be retained to the maximum extent, and high-quality infrared to visible light image translation can be achieved. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] Figure 1 It is a flow chart of an infrared to visible light image translation method based on contrast learning of the present invention;

[0037] Figure 2 is a schematic diagram of an encoder in the improved ViT model of the present invention;

[0038] In the figure, (a) - schematic diagram of a single coding block, (b) - schematic diagram of the encoder processing process;

[0039] Figure 3 It is a schematic diagram of single-layer contrast loss;

[0040] Figure 4 It is a multi-scale discriminator model diagram;

[0041] Figure 5 This is an example of the infrared to visible light image translation result. DETAILED DESCRIPTION

[0042] The present invention is further described below in conjunction with the accompanying drawings and embodiments. The present invention includes but is not limited to the following embodiments.

[0043] The present invention provides an infrared to visible light image translation method based on contrastive learning, such as Figure 1 As shown, the specific implementation process is as follows:

[0044] 1. Build an image dataset

[0045] Image datasets are used to train and test image translation models, including infrared images and visible light images. The present invention uses InfraredCity-Lite and IRVI, datasets dedicated to infrared-visible light image or video translation, for training and testing (reference: Yu Z, Chen K, Li S, et al. ROMA: Cross-Domain Region Similarity Matching for Unpaired Nighttime Infrared to Daytime Visible Video Translation [C] / / Proceedings of the 30th ACM International Conference on Multimedia. Lisboa Portugal: ACM, 2022: 5294-5302.). InfraredCity-Lite consists of nighttime infrared video data and daytime visible light video data, and is divided into two major parts: Traffic and Monitor. It contains 37,339 video training frames and 4500 test frames, with an image resolution of 256*256, and is used to convert nighttime infrared videos into visible light videos. IRVI provides a dataset for daytime infrared to visible light translation.

[0046] 2. Build an image translation model

[0047] The details are as follows:

[0048] Step 2-1, improved ViT feature extraction module. Use the ViT model pre-trained on a large-scale dataset to extract features and obtain token embeddings to prepare for the subsequent calculation of the loss function. The specific implementation is as follows:

[0049] Given a two-dimensional image , the ViT model first divides the image into N patches, and then converts it into a vector sequence, which also contains category embedding and position embedding. Then the image vector sequence is sent to the Transformer Encoder, as shown in Figure 2 As shown in (a), the encoder consists of L encoder blocks, each of which consists of a multi-head attention sublayer and an MLP sublayer. In the figure, Embedded Patches represents the input image vector sequence, and Norm represents the normalization calculation. In the present invention, an encoder structure with L=12 is used to feed the image sequence into it. The classic ViT model will directly use the tag embedding finally output by the encoder, feed it into an MLP layer for classification tasks, and then finally output only the category embedding vector. The present invention hopes to make full use of the tag embedding extracted by the encoding blocks of different layers for subsequent tasks. Therefore, Figure 2 As shown in (b), during the encoding process, the n-layer tag embeddings are first directly extracted, and then their category embeddings are deleted as the final n-layer features. Each layer of features is a sequence with a length of N and a dimension of D. N represents the number of blocks into which the image is divided, and the range of possible values ​​depends on the image size. D represents the sequence dimension, and the common value is 512. n represents that the selected features come from the n-th layer coding block, and the range of possible values ​​is (1,6). Figure 2 In (b), the label embeddings of the 1st, 3rd, and 5th layer encoding blocks are used as output features.

[0050] Step 2-2, improved multi-layer contrast loss. Single-layer contrast loss is as follows Figure 3 As shown, for the input image x∈X and its corresponding generated image Send them into the ViT model respectively to obtain the labeled embedding features of a total of I layers. The i-th layer labeled embedding features are expressed as i∈[1,I], let the ViT model be V, which is represented as follows:

[0051]

[0052] Where X represents the input image dataset, Represents the generation of an image dataset.

[0053] for and , sample them at the same position to make their sequence length and dimension become K, and normalize the sampled sequence vector to a unit ball, and get i∈[1,I]; in Take a vector For query, Take the vector at the corresponding position is positive, the remaining vectors is negative, for z,z + and z - Apply the infoNCE loss function:

[0054]

[0055] Among them, l(z,z + ,z - ) represents z,z + and z - The infoNCE loss function value of the application, τ represents the temperature coefficient, and the value is τ=0.07; the common value of K is 256.

[0056] Embed the i-th layer tag ,use Denotes the jth vector, j∈[1,K], which embeds the i-th layer tag ,use represents the jth vector, and the remaining vectors are represented as , the improved multi-layer contrast loss is calculated as follows:

[0057]

[0058] Among them, L MNCE Represents the multi-layer contrast loss value, Represents the mathematical expectation of a computation.

[0059] Step 2-3, the generative adversarial network consisting of the ResNet generator and the multi-scale discriminator.

[0060] The ResNet generator is a residual generator that takes in an input image and outputs a generated image.

[0061] The multi-scale discriminator determines the authenticity of the generated image, such as Figure 4 As shown, the specific implementation is as follows:

[0062] First, the real image y and the generated image They are sent to the ViT model to obtain a layer of labeled embedding features T y , , and then reshape the label embedding into the same three-dimensional form set P as the image y , ,in Then, non-overlapping spatial averages are performed on adjacent vectors with receptive fields of different sizes, including P y , After spatial averaging with a receptive field of 2×2, we get the set P y 2 , , and then through the spatial average of the receptive field of 4×4, we get the set P y 4 , , and so on until all spatial average sets of different scales are obtained; then, all the obtained sets are restored to the form of labeled embedding, including P y 2 , Restore to T y 2 , , P y 4 , Restore to T y 4 , , so that all the recovered sets are obtained, and they are concatenated to obtain and ; Finally, and They are sent to a three-layer learnable MLP classification network respectively, and the MLP outputs the result of whether it is true or false. Figure 4 In the figure, token embeddings represents the token embedding feature, reshape represents the size transformation process, average represents the spatial average process, Concatenate represents the concatenation process, Fake represents fake, and Real represents real.

[0063] 3. Training Generative Adversarial Networks

[0064] The purpose of training a generative adversarial network is to learn the cross-domain mapping between the two infrared image domains and the visible light image domain to complete the image translation task. This includes: sending the infrared image to the generator to generate the corresponding visible light image, and then sending the generated image to the discriminator to judge the authenticity of the generated image, and using adversarial loss and multi-layer contrast loss to train the generator and discriminator. The details are as follows:

[0065] Given an infrared image dataset and visible light image datasets , the generator G learns the mapping relationship from X to Y, converts X into an image with the features of Y, but still maintains its original structure, and obtains the generated image , the discriminator D distinguishes the image and the real image sample y∈Y to output the result of judging the authenticity of the generated image, where G attempts to generate images G(x) that look similar to images from domain Y, while D aims to distinguish between the translated samples G(x) and the real samples y.

[0066] The loss function is used to train the entire generative adversarial network, and the network is trained by minimizing the loss function of the discriminator and the loss function of the generator. The loss function of the discriminator L D for:

[0067]

[0068] Among them, L D represents the loss function value of the discriminator, G represents the generator, and D represents the multi-scale discriminator. MPL represents the MLP classification network, Indicates the real image feature set obtained by the intermediate calculation of the discriminator, Represents the set of generated image features obtained by the intermediate calculation of the discriminator.

[0069] The loss function of the generator is:

[0070]

[0071] Among them, L G represents the loss function value of the generator, is an adversarial loss to encourage the generator to generate images that are visually similar to the target domain, L MNCE is the multi-layer contrast loss of the present invention, which is used to maintain the spatial consistency between the input image and the output image, and λ is the control L MNCE The weight of is in the range of (0,1).

[0072] 4. Image Translation

[0073] The infrared image to be translated is sent to the generator of the trained adversarial generative network to obtain the final visible light image.

[0074] Examples of infrared to visible light image translation results are as follows Figure 5As shown in the figure, the first column Nighttime Infrared is a nighttime infrared image, wherein the first and second rows of images are nighttime infrared images taken on urban roads, and the third and fourth rows of images are nighttime infrared images taken on highways; the second column VCT-GAN is the result obtained by translating the first column of nighttime infrared images using the method of the present invention (VCT-GAN); the third column EADS is the result obtained by translating the first column of nighttime infrared images using the EADS method; the fourth column CUT is the result obtained by translating the first column of nighttime infrared images using the CUT method; the fifth column I2V-GAN is the result obtained by translating the first column of nighttime infrared images using the I2V-GAN method; the sixth column ROMA is the result obtained by translating the first column of nighttime infrared images using the ROMA method.

[0075] Since the method of the present invention (VCT-GAN) innovatively combines the feature extraction module of ViT with the classic image translation algorithm CUT (Contrastive Learning for Unpaired Image-to-Image Translation), an improved multi-layer contrast loss function is constructed, which better maintains the spatial consistency in image translation, retains the scene structure information to the maximum extent, and can achieve high-quality infrared to visible light image translation.

Claims

1. A method for infrared to visible light image translation based on contrastive learning, characterized in that Here are the steps: Step 1: construct an image dataset for training and testing an image translation model, wherein the image dataset includes infrared images and visible light images; Step 2: construct an image translation model, the image translation model includes: Improved ViT feature extraction module for image feature extraction and obtaining marker embedding; Improved multi-layer contrast loss, used to calculate the mutual information of input and output images, using its maximum correlation as one of the loss functions; The generative adversarial network composed of a ResNet generator and a multi-scale discriminator takes an infrared image as input and generates the corresponding visible light image. Step 3, training the generative adversarial network, specifically comprising: sending the infrared image to the generator to generate a corresponding visible light image, and then sending the generated image to the discriminator to judge whether the generated image is true or false, and using adversarial loss and multi-layer contrast loss to train the generator and the discriminator; Step 4: Input the infrared image into the trained generator to obtain a visible light image and complete the image translation.

2. The infrared to visible light image translation method based on contrastive learning as claimed in claim 1, characterized in that: In step 1, the InfraredCity-Lite and IRVI datasets dedicated to infrared-visible light image or video translation are used for model training and testing.

3. The infrared to visible light image translation method based on contrastive learning as claimed in claim 1, characterized in that: In step 2, an image translation model is constructed as follows: Step 2-1, improved ViT feature extraction module: use the ViT model pre-trained on a large-scale dataset to perform feature extraction and obtain tag embedding. Specifically: given a two-dimensional image, the ViT model first divides the image into N patches, and then converts it into a vector sequence, which also includes category embedding and position embedding; then the image vector sequence is sent to the Transformer encoder. During the encoding process, the tag embedding of the n layers is first extracted, and then its category embedding is deleted as the final output of the n-layer feature, each layer of feature is a sequence with a length of N and a dimension of D; the encoder is composed of L encoding blocks, each encoding block is composed of a multi-head attention sublayer and an MLP sublayer, and L is set to 12; the D represents the sequence dimension, which is set to 512, and the n is an integer in the range of (1,6); Step 2-2, improved multi-layer contrast loss: For the input image x∈X and its corresponding generated image Send them into the ViT model respectively to obtain the labeled embedding features of a total of I layers. The i-th layer labeled embedding features are expressed as The ViT model is denoted as V, which is expressed as follows: Where X represents the input image dataset, Indicates the generation of image dataset; for and Sampling is performed at the same position to make the sequence length and dimension become K, and the sampled sequence vector is normalized to a unit ball to obtain exist Take a vector For query, Take the vector at the corresponding position is positive, the remaining vectors is negative, for z,z + and z - Apply the infoNCE loss function: Among them, l(z,z + ,z - ) represents z,z + and z - The infoNCE loss function value of the application, τ represents the temperature coefficient, and the value is τ=0.07; the K is set to 256; Embed the i-th layer tag use Denotes the jth vector, j∈[1,K], which embeds the i-th layer tag use represents the jth vector, and the remaining vectors are represented as The improved multi-layer contrast loss is calculated as follows: Among them, L MNCE Represents the multi-layer contrast loss value, Represents the mathematical expectation of a calculation; Step 2-3, the generative adversarial network composed of the ResNet generator and the multi-scale discriminator: the ResNet generator is a residual generator, which inputs the input image and outputs the generated image; the multi-scale discriminator judges the authenticity of the generated image, as follows: First, the real image y and the generated image Send them into the ViT model to get a layer of labeled embedding features Then reshape the labeled embedding into a three-dimensional set like the image in Then, non-overlapping spatial averages are performed on adjacent vectors with receptive fields of different sizes, including After spatial averaging with a receptive field of 2×2, the set After the spatial average with a receptive field of 4×4, the set This process is continued until all spatial average sets of different scales are obtained; then, all the obtained sets are restored to the form of labeled embedding, including Restore to Will Restore to In this way, all the recovered sets are obtained, and they are spliced ​​together to obtain and Finally, and They are sent to a three-layer learnable MLP classification network respectively, and the MLP outputs the result of whether it is true or false.

4. The infrared to visible light image translation method based on contrastive learning according to claim 1 or 3, characterized in that: In step 3, the generative adversarial network is trained as follows: Given an infrared image dataset and visible light image datasets The generator G learns the mapping relationship from X to Y, converts X into an image with the features of Y, but still maintains its original structure, and obtains the generated image The discriminator D distinguishes images And the real image sample y∈Y to output the result of judging whether the generated image is true or false; The loss function is used to train the entire generative adversarial network, and the network is trained by minimizing the loss function of the discriminator and the loss function of the generator. The loss function of the discriminator L D for: Among them, L D represents the loss function value of the discriminator, G represents the generator, and D represents the multi-scale discriminator. MPL represents the MLP classification network, Indicates the real image feature set obtained by the intermediate calculation of the discriminator, Represents the generated image feature set obtained by the intermediate calculation of the discriminator; The loss function of the generator is: Among them, L G represents the loss function value of the generator, is the adversarial loss, L MNCE is the multi-layer contrast loss, λ is the control L MNCE The weight of is in the range of (0,1).

Images