Color migration optimization method based on generative adversarial network

Abstract

The invention discloses a color migration optimization method based on a generative adversarial network, and relates to the technical field of image processing, and the method comprises the following steps: data preparation: collecting a degraded image and a corresponding high-quality real image, constructing a training data set, preprocessing the image, and dividing the image into a training set and a test set; a U-Net structure is adopted as a generator main body, an encoder extracts fade video frame features, a decoder recovers image details, and meanwhile an attention mechanism is introduced to pay attention to key color areas; and discriminator construction: constructing a discriminator by using a multilayer convolutional neural network. According to the color migration optimization method based on the generative adversarial network, image features can be more effectively extracted and utilized by adopting a generator of a U-Net structure and an attention mechanism, more accurate color migration is realized, and the quality and naturalness of a generated image are improved; and the detail recovery capability is enhanced.

Description

Technical Field

[0001] The present invention relates to the field of image processing technology, and in particular to a color migration optimization method based on a generative adversarial network. Background Art

[0002] In the field of image processing, color transfer technology aims to transfer the color style of one image to another, thereby optimizing and enhancing the image's colors. Traditional color transfer methods, primarily based on statistical analysis and handcrafted features, often suffer from poor color transfer results and loss of image detail when processing complex images and diverse color requirements.

[0003] With the development of deep learning technology, image generation and translation methods based on generative adversarial networks (GANs) have made significant progress. Through adversarial training of the generator and discriminator, GANs can learn the distribution of real images and generate highly realistic images. However, existing GAN-based color transfer methods still need to be improved in terms of image detail restoration, color naturalness, and training stability. Summary of the Invention

[0004] The purpose of the present invention is to provide a color transfer optimization method based on a generative adversarial network to solve the problems raised in the above background technology.

[0005] To achieve the above objectives, the present invention provides the following technical solution: a color transfer optimization method based on a generative adversarial network, comprising the following steps: data preparation: collecting degraded images and their corresponding high-quality real images, constructing a training dataset, preprocessing the images and dividing them into a training set and a test set;

[0006] Generator Design: The U-Net structure is used as the main generator. The encoder extracts features of faded video frames, and the decoder restores image details. At the same time, an attention mechanism is introduced to focus on key color areas.

[0007] Discriminator construction: Use a multi-layer convolutional neural network to build a discriminator, extract features from the input image and output the probability that it is a real image;

[0008] Adversarial training: The generator generates color-enhanced images, and the discriminator distinguishes between real and generated images. Through alternating training, the generator parameters are optimized using adversarial loss combined with pixel-level loss.

[0009] Inference application: Use the trained generator to optimize color transfer of images or video frames.

[0010] Furthermore, in the data preparation step, image preprocessing includes image cropping and normalization, which uniformly adjusts the images to the same size and maps the pixel values to the [-1, 1] interval.

[0011] Furthermore, the convolutional layer of the U-Net encoder uses a 3×3 convolution kernel with a stride of 1 and a padding of 1, and the downsampling layer uses a maximum pooling operation; the decoder uses transposed convolution and deconvolution operations, and the corresponding layers of the encoder and decoder perform feature fusion.

[0012] Furthermore, the attention mechanism introduced in the generator obtains the global information of the feature map through global average pooling, calculates the attention weight through the multi-layer perceptron, and performs weighted processing on the feature map.

[0013] Furthermore, in the multi-layer convolutional neural network of the discriminator, the convolution layer uses a 3×3 convolution kernel, a step size of 1, and a padding of 1. As the network depth increases, the number of convolutional layer channels increases, and finally the probability value is output through the fully connected layer and the sigmoid activation function.

[0014] Furthermore, in adversarial training, the training data of the generator and discriminator are divided into batches, each batch contains 32 pairs of images, and the Adam optimizer is used to update the network parameters and train for 200 epochs.

[0015] Furthermore, in adversarial training, the pixel-level loss of the generator adopts mean square error loss.

[0016] Furthermore, in inference applications, when optimizing the color migration of a video, the video frames are processed frame by frame and reassembled into a video.

[0017] The present invention provides a color migration optimization method based on a generative adversarial network, which has the following beneficial effects: by adopting a U-Net structure and a generator with an attention mechanism, the present invention can more effectively extract and utilize image features, achieve more accurate color migration, and improve the quality and naturalness of the generated image; enhance the ability to restore details: the encoder and decoder structure and feature fusion method of U-Net, combined with the attention mechanism's focus on key areas, enable the generated image to better restore detail information and reduce the problem of detail loss; adopt an adversarial loss combined with a pixel-level loss method to train the generator, and use a multi-layer convolutional neural network as a discriminator, so that the model training is more stable and can effectively avoid problems such as mode collapse during training. DETAILED DESCRIPTION

[0018] The following embodiments of the present invention are further described in detail with reference to the examples. The following examples are used to illustrate the present invention but are not intended to limit the scope of the present invention. A color migration optimization method based on a generative adversarial network comprises the following steps: Data preparation stage:

[0019] Data Collection: Collect a large number of degraded images and their corresponding high-quality real-world images from public image databases, internet resources, and real-world application scenarios (such as historical imagery and old photo libraries). Ensure that the collected data covers a variety of scenes, color styles, and image content.

[0020] Data preprocessing: The collected images are cropped and resized to a suitable size, such as 256×256 pixels. The images are then normalized to map pixel values to the range [-1, 1] to meet the input requirements of the neural network.

[0021] Dataset partitioning: The preprocessed dataset is divided into a training set and a test set in a ratio of 8:2. The training set is used for model training, and the test set is used to evaluate model performance.

[0022] Generator design and construction:

[0023] U-Net architecture: Use a deep learning framework (such as PyTorch or TensorFlow) to build a U-Net generator. In the encoder, construct multiple convolutional and downsampling layers in sequence. Each convolutional layer uses a 3×3 kernel, a stride of 1, and padding of 1. The downsampling layer uses max pooling with a 2×2 kernel size and a stride of 2. In the decoder, use upsampling layers (such as transposed convolution) and deconvolution layers to restore image size and reconstruct features, again using appropriate convolution kernels and stride settings.

[0024] Attention Mechanism: Add an attention mechanism module to key layers of the generator's encoder and decoder (such as intermediate layers and those near the output layer). This module calculates attention weights for feature maps, weighting features from different regions. In practice, global average pooling is used to capture global information from the feature map. Attention weights are then calculated using a multi-layer perceptron (MLP). Finally, the attention weights are multiplied and added to the original feature map to produce an enhanced feature map.

[0025] Discriminator construction:

[0026] Convolutional Neural Network Construction: A multi-layer convolutional neural network is constructed as the discriminator. The discriminator takes image data as input and first extracts features through multiple convolutional layers. The convolutional layers also use a 3×3 convolution kernel with a stride of 1 and padding of 1. As the network progresses, the number of channels in the convolutional layers is gradually increased to extract more advanced image features.

[0027] Fully connected layer and output: After passing through multiple convolutional layers, the extracted feature maps are flattened and then connected to the fully connected layer. The number of neurons in the fully connected layer is set according to the actual situation. Finally, a sigmoid activation function is used to output a probability value representing the probability that the input image is a real image.

[0028] Adversarial Training:

[0029] Initialize the model and optimizer: Initialize the network parameters of the generator and discriminator using random initialization or pre-trained model initialization. Select the Adam optimizer for the generator and discriminator, and set hyperparameters such as the learning rate, β1, and β2. For example, set the learning rate to 0.0002, β1 to 0.5, and β2 to 0.999.

[0030] Training Process: In each round of training, a batch of 32 pairs of degraded and real images is extracted from the training set. The degraded images are fed into the generator to produce color-optimized images. The real and generated images are fed into the discriminator to calculate the discriminator output. Gradients are calculated based on the loss functions of the generator and discriminator (an adversarial loss combined with a pixel-level loss, where the mean squared error loss can be used), and the Adam optimizer is used to update the network parameters of the generator and discriminator.

[0031] Training Monitoring and Convergence Assessment: During training, the loss values of the generator and discriminator are recorded for each round of training. By observing the changing trends of the loss values, the model is considered to have converged when the loss values of the generator and discriminator stabilize and the quality of the generated images no longer significantly improves. Training is stopped. The entire training process lasts for 200 epochs and takes approximately 12 hours.

[0032] Reasoning Application:

[0033] Model loading: Load the trained generator model. In actual applications, the trained model can be saved as a file and loaded when needed.

[0034] Image or video processing: For a single image that needs to be optimized for color migration, it is pre-processed (such as normalization) and then input into the generator, and the generator outputs the color-optimized image. For video processing, the video frames are read frame by frame, pre-processed and then input into the generator, and the generated optimized frames are saved and reassembled into a video. The above is only a specific embodiment of the present invention, but the protection scope of the present invention is not limited to this. Any technician familiar with this technical field can easily think of changes or replacements within the technical scope disclosed by the present invention, which should be covered within the protection scope of the present invention. Therefore, the protection scope of the present invention shall be based on the protection scope of the claims.

Claims

1. A color transfer optimization method based on generative adversarial network, characterized in that: The following steps are included: Data preparation: collect degraded images and their corresponding high-quality real images, build a training dataset, preprocess the images and divide them into training and test sets; Generator Design: The U-Net structure is used as the main generator. The encoder extracts features of faded video frames, and the decoder restores image details. At the same time, an attention mechanism is introduced to focus on key color areas. Discriminator construction: Use a multi-layer convolutional neural network to build a discriminator, extract features from the input image and output the probability that it is a real image; Adversarial training: The generator generates color-enhanced images, and the discriminator distinguishes between real and generated images. Through alternating training, the generator parameters are optimized using adversarial loss combined with pixel-level loss. Inference application: Use the trained generator to optimize color transfer of images or video frames.

2. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: In the data preparation step, image preprocessing includes image cropping and normalization, which uniformly adjusts the images to the same size and maps the pixel values to the [-1, 1] interval.

3. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: The convolutional layer of the U-Net encoder uses a 3×3 convolution kernel with a stride of 1 and a padding of 1, and the downsampling layer uses a maximum pooling operation; the decoder uses transposed convolution and deconvolution operations, and the corresponding layers of the encoder and decoder perform feature fusion.

4. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: The attention mechanism introduced in the generator obtains the global information of the feature map through global average pooling, calculates the attention weight through the multi-layer perceptron, and performs weighted processing on the feature map.

5. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: In the multi-layer convolutional neural network of the discriminator, the convolution layer uses a 3×3 convolution kernel, a step size of 1, and a padding of 1. As the network depth increases, the number of convolutional layer channels increases, and finally the probability value is output through the fully connected layer and the sigmoid activation function.

6. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: In adversarial training, the training data of the generator and discriminator are divided into batches, each batch contains 32 pairs of images, and the Adam optimizer is used to update the network parameters and train for 200 epochs.

7. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: In adversarial training, the pixel-level loss of the generator adopts mean square error loss.

8. The color transfer optimization method based on generative adversarial network according to claim 1, characterized in that: In inference applications, when optimizing color migration for a video, the video frames are processed frame by frame and reassembled into a video.