Real-time compensation method and device for color cast and distortion of endoscopic images

Abstract

The method and device for real-time compensation of color cast and distortion of endoscopic images can compress the GAN generator and ensure the generation details of large-sized endoscopic images, and improve the quality of images generated by the compression model without adding additional generator structure. The method uses the GAN compression based network as the standard network, which includes: (1) pre-training a cycle GAN network with a multi-scale discriminator added to obtain a teacher generator and a discriminator; (2) compressing the teacher generator using knowledge distillation and neural structure search, and inheriting the multi-scale discriminator to train a student generator; (3) evaluating and selecting the appropriate compression model structure to obtain the final generator model.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image processing, and in particular to a method for real-time compensation of color cast and distortion of endoscopic images, and a device for real-time compensation of color cast and distortion of endoscopic images. Background Art

[0002] Photodynamic surgery is a rapidly developing therapy in recent years. This therapy uses laser, oxygen, and photosensitizer to react with lesions to kill tissue cells. Its unique mechanism of action has become one of the research hotspots in the field of tumor prevention and treatment. When performing photodynamic surgery with the help of endoscopes in vivo, turning on the laser causes blindness in the endoscope field, which seriously affects the doctor's operation. By placing a filter in front of the endoscope lens charge-coupled device (CCD), the light in the 630-nanometer band is shielded to eliminate the overexposure effect, but at the same time, the color distortion of the endoscope image is introduced. Comparison of endoscope images under three different lighting conditions: normal lighting scene, laser blindness scene, and laser scene after adding an endoscope filter. It can be seen that the image with the filter added has obvious color difference, and the color texture of tissue organs is an important basis for doctors' diagnosis. Due to the real-time requirements of surgical operations, it is of great significance to quickly correct the color cast image during surgery to improve the accuracy of surgical operations.

[0003] In recent years, deep learning methods have not only achieved remarkable results in the fields of image restoration and color correction, but the development of model compression algorithms has also improved the model operation speed, making it possible to apply deep learning methods to endoscopic surgery color restoration scenarios. However, when the compression model is actually applied, there is a problem that the processed endoscopic image quality becomes blurred.

[0004] In recent years, some progress has been made in the field of image color processing. A. Galdran et al. proposed a method for short-wave color restoration of underwater images based on the red channel, and the polynomial regression color correction method has also found new application areas. Since Z. Cheng et al. first introduced deep learning into the field of image colorization, deep convolutional neural networks (DCNNs) have rapidly developed into the main application technology in this field, with impressive results in image color style transfer, color correction, and image colorization. J. Li et al. proposed WateGAN to achieve underwater image modeling and color correction from air RGB-D images, and JW Su et al. proposed Instance-aware Image Colorization to achieve multi-target colorization and fusion. However, for the field of medical image processing, endoscope datasets with paired matching (pixels and regions) are very difficult to obtain. The emergence of cycle GAN has provided powerful tools for image conversion of unpaired data, such as CT and MRI image conversion, and endoscope image generation tasks.

[0005] A major problem that restricts the practical application of the DCNN method for endoscopic image color restoration is the model reasoning speed. In electronic endoscope-assisted photodynamic therapy surgery, doctors need to operate surgical instruments based on processed images, which requires the model to process large-size images at a speed that can meet the requirements of real-time surgery. Model compression can effectively improve the speed of model reasoning. Since S. Hanet al. summarized the methods in the field of model compression, this research field has become one of the hot development directions of deep learning, and many methods have emerged, including model pruning, lightweight module design, knowledge distillation, low-rank decomposition, etc. These methods have achieved good results in image classification task models, but due to the instability of GAN network training and other reasons, the above methods are not competent for GAN generator compression tasks. Knowledge distillation methods are popular methods in the field of model compression. In classification tasks, the target distillation method has been proven to be effective in using the category probability output from softmax as a learning target. In GAN tasks, feature distillation is used. On the one hand, since the GAN task outputs a certain image rather than a probability distribution, it is difficult to extract key knowledge from the output layer; on the other hand, the number of channels in the intermediate layer is more, and the knowledge provided is richer. Recently, research on compression of GAN models has achieved remarkable results, such as the GAN Slimming method proposed by H. Wang et al. and the GAN compression method proposed by M. Li et al. However, experiments have found that in the task of color restoration of 1920*1080 endoscopic images, both methods have the problem of image blurring and loss of details. Therefore, a training strategy that can compress the GAN generator and ensure the generation of details of larger endoscopic images is of great significance. Summary of the invention

[0006] In order to overcome the defects of the prior art, the technical problem to be solved by the present invention is to provide a real-time compensation method for color cast and distortion of endoscopic images, which can compress the GAN generator and ensure the generation details of larger-sized endoscopic images, thereby improving the quality of images generated by the compression model without adding any additional generator structure.

[0007] The technical solution of the present invention is: this method for real-time compensation of color cast and distortion of endoscopic images uses a generative adversarial compression network (GAN compression based network) as a standard network, which includes the following steps:

[0008] (1) Pre-train a cycle-consistent generative adversarial network (cycle GAN) with a multi-scale discriminator to obtain the teacher generator and discriminator;

[0009] (2) The teacher generator is compressed using knowledge distillation and neural structure search, and the student generator is trained by inheriting the multi-scale discriminator;

[0010] (3) Evaluate and screen the appropriate compression model structure to obtain the final generator model.

[0011] The present invention uses a generative adversarial compression network (GAN compression based network) as the standard network, pre-trains a cycle-consistent generative adversarial network with a multi-scale discriminator to obtain a teacher generator and a discriminator, compresses the teacher generator using knowledge distillation and neural structure search, and inherits the multi-scale discriminator to train a student generator, evaluates and screens the appropriate compression model structure to obtain the final generator model, thereby being able to compress the GAN generator and ensure the generation details of larger-sized endoscopic images, thereby improving the quality of images generated by the compression model without adding an additional generator structure.

[0012] A real-time compensation device for color cast and distortion of endoscopic images is also provided, which uses a generative adversarial compression network as a standard network and includes:

[0013] A pre-training module, which is configured to pre-train a cycle-consistent generative adversarial network with a multi-scale discriminator added to obtain a teacher generator and discriminator;

[0014] The training module is configured to compress the teacher generator using knowledge distillation and neural structure search, and inherit the multi-scale discriminator to train the student generator;

[0015] A screening module is configured to evaluate and screen the appropriate compression model structure to obtain the final generator model. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] Figure 1 The figure is a flow chart of the method for real-time compensation of color cast and distortion of endoscopic images according to the present invention. DETAILED DESCRIPTION

[0017] like Figure 1 As shown, this method for real-time compensation of color cast distortion of endoscopic images uses a generative adversarial compression network as a standard network, which includes:

[0018] (1) Pre-train a cycle-consistent generative adversarial network with a multi-scale discriminator to obtain the teacher generator and discriminator;

[0019] (2) The teacher generator is compressed using knowledge distillation and neural structure search, and the student generator is trained by inheriting the multi-scale discriminator;

[0020] (3) Evaluate and screen the appropriate compression model structure to obtain the final generator model.

[0021] The present invention uses a generative adversarial compression network as a standard network, pre-trains a cycle-consistent generative adversarial network with a multi-scale discriminator added to obtain a teacher generator and a discriminator, compresses the teacher generator using knowledge distillation and neuron structure search, and inherits the multi-scale discriminator to train a student generator, evaluates and screens the appropriate compression model structure to obtain the final generator model, thereby being able to compress the GAN generator and ensure the generation details of larger-sized endoscopic images, thereby improving the quality of images generated by the compression model without adding an additional generator structure.

[0022] Preferably, in step (1), the input real sample and fake sample images are downsampled twice, and then the small-scale images and the original-size images are respectively input into two discriminators D with the same structure. 1 and D 2 , respectively judging the authenticity of the image from different scales. The image is considered to be true only when the judgment results of both discriminators are true. If the output results of both discriminators are false or one is true and the other is false, the image is judged to be false.

[0023] Preferably, in step (1), D 1 The receptive field is smaller, guiding the generator to generate more detailed images, D 2 The generator is guided by a larger receptive field to generate a globally consistent image. At this time, the optimization objective function of the generative adversarial network is formula (1):

[0024]

[0025] Among them, L GAN is the adversarial loss, and the training goal of the generator G is to make L GAN As small as possible, the discriminator D i The optimization goal is to make L GAN As big as possible.

[0026] Preferably, in step (2), during compression, no matter what the teacher generator generates, its image is selected as a positive sample to guide the training of the student generator; for a given input data I filter , the data generated by the teacher generator is represented as I t , the data generated by the student generator is I s , the training images at this time have one-to-one correspondence in structure and texture details; the training strategy is expressed as:

[0027]

[0028] in, ‖ ‖ 1 is L1 regularization.

[0029] Preferably, in step (2), a multi-scale discriminator is used, and the generative adversarial network optimization objective function (GAN objective) is expressed in combination with formula (1) as follows:

[0030]

[0031] The discriminator of the student network has the same structure as the discriminator of the teacher network. During training, the weights of the multi-scale discriminator of the teacher network are used to initialize the student discriminator. The weights are not fixed but are continuously adjusted as the compression model is trained.

[0032] Preferably, in step (2), feature extraction is performed from the middle layer of the teacher generator. and are the nth intermediate layers to be distilled in the teacher model and the student model, respectively. N represents the total number of layers. The distillation loss is expressed as:

[0033]

[0034] Since the number of output channels of the corresponding layers of the teacher model and the student model is different, a feature mapping function F is required before calculation n The output of the intermediate layer of the student model is transformed into a space with the same dimension as the teacher model. The mapping consists of a 1×1 learnable convolutional layer. The distillation process is to optimize and F n To minimize L distill .

[0035] Preferably, in step (2), when training the student model, a super-net containing multiple candidate channel numbers is first trained, and different channel numbers are arranged and combined to form sub-nets with different structures. Sub-nets and super-net parameters are shared using a once-for-all structure, and sub-nets obtain corresponding layer weights from super-net; the channels of the nth layer excluding the input and output layers of the student network are represented as [C n ], set subnets[C n ] The optional channel number range is [16, 24, 32, 48, 64], each layer has 5 channel number structures, and the entire candidate network pool has 5 N Different sub-network structures are randomly selected from the sub-network pool during training, and forward propagation calculation is performed with a certain number of channels, and weights are updated by back propagation according to formula (6); when searching for the optimal network, an evolutionary algorithm is used to screen the sub-network structure.

[0036] L=LGAN +λ recon L recon +λ distill L distill +λ FM L FM (6)

[0037] Among them, λ recon ,λ distill ,λ FM These are all hyperparameters.

[0038] Preferably, in step (2), Fréchet Inception Distance and Multiply-Accumulate Operations are used as evaluation criteria, and two samples are randomly selected from the candidate pool each time to evaluate FID and MACs, and candidates with both high values ​​are eliminated, and candidates with lower values ​​are retained as parents, which are fine-tuned to generate offspring and put into the candidate pool until the screening is completed.

[0039] Preferably, in step (3), in order to make the student generator G s The training process is more stable, and it can learn image features of different scales at the same time. It uses the feature matching loss function based on the discriminator to extract features from the non-output layer of the discriminator, so that the student model can learn the data generated by the teacher generator. t and data generated by the student generator I s Matching of features between; defining sub-discriminator D i The feature extractor of the jth layer is Then the feature matching loss function L FM It is expressed as:

[0040]

[0041] Where M is the total number of layers, N j Indicates the number of elements in each layer, D i Not participating in L FM The optimization process only acts as a feature extractor.

[0042] A person skilled in the art can understand that all or part of the steps in the above-mentioned embodiment method can be completed by instructing the relevant hardware through a program, and the program can be stored in a computer-readable storage medium. When the program is executed, it includes the steps of the above-mentioned embodiment method, and the storage medium can be: ROM / RAM, disk, CD, memory card, etc. Therefore, corresponding to the method of the present invention, the present invention also includes a real-time compensation device for color cast and distortion of endoscopic images, which is usually represented in the form of functional modules corresponding to the steps of the method. The device is a real-time compensation device for color cast and distortion of endoscopic images, and the device uses a GAN compression based network as a standard network, which includes:

[0043] A pre-training module, which is configured to pre-train a cycle-consistent generative adversarial network with a multi-scale discriminator added to obtain a teacher generator and discriminator;

[0044] The training module is configured to compress the teacher generator using knowledge distillation and neural structure search, and inherit the multi-scale discriminator to train the student generator;

[0045] A screening module is configured to evaluate and screen the appropriate compression model structure to obtain the final generator model.

[0046] For the model compression strategy adopted by the present invention, the performance of the student model is improved by improving the discriminator instead of increasing the network depth recompression of the teacher generator for the following reasons:

[0047] 1. In order to ensure that the model processing speed is not affected, the image quality cannot be improved by deepening the generator network or adding other structures, because larger model parameters will reduce the prediction speed and affect the doctor's viewing and operation during surgery.

[0048] 2. Improving the accuracy of the teacher discriminator and then compressing it may not necessarily achieve the desired effect. The work of Cho and Hariharan showed that mismatched capacity makes the student network unable to stably imitate the teacher network. Improving the accuracy of the teacher network does not mean that the ability of the student model is improved.

[0049] 3. The features learned by the discriminator help improve the output performance of the model. The Metropolis-Hastings GAN study shows that the discriminator contains some useful information that helps the generator further improve the output results, so it is worth trying to use the improved teacher discriminator to improve the compression model training strategy of the student generator. In order to solve the problem of blurry output images of the compression model, this paper chooses to introduce a multi-scale discriminator for improvement, and uses a feature matching loss function to make the compression model training process more stable.

[0050] The above description is only a preferred embodiment of the present invention and does not limit the present invention in any form. Any simple modification, equivalent changes and modifications made to the above embodiments based on the technical essence of the present invention are still within the protection scope of the technical solution of the present invention.

Claims

1. A real-time compensation method for color cast and distortion of endoscopic images, characterized by: The method uses a generative adversarial compression network as a standard network, which includes the following steps: (1) Pre-train a cycle-consistent generative adversarial network with a multi-scale discriminator to obtain the teacher generator and discriminator; (2) The teacher generator is compressed using knowledge distillation and neural structure search, and the student generator is trained by inheriting the multi-scale discriminator; (3) Evaluate and select the appropriate compression model structure to obtain the final generator model; In step (3), in order to make the student generator G s The training process is more stable, and it can learn image features of different scales at the same time, using the discriminator-based feature matching loss function L FM ; In the step (1), the input real sample and fake sample images are downsampled twice, and then the small-scale images and the original-size images are respectively input into two discriminators D1 and D2 with the same structure, and the authenticity of the images is judged from different scales. Only when the judgment results of the two discriminators are both true, the image is considered to be true. If the output results of the two discriminators are both false or one is true and the other is false, the image is judged to be false. In step (1), D1 has a smaller receptive field and guides the generator to generate a more refined image. D2 guides the generator to generate a globally consistent image from a larger receptive field. At this time, the optimization objective function of the generative adversarial network is formula (1): Among them, L GAN is the adversarial loss, and the training goal of the generator G is to make L GAN As small as possible, the discriminator D i The optimization goal is to make L GAN As big as possible; In step (2), during compression, no matter what the teacher generator generates, its image is selected as a positive sample to guide the training of the student generator; for a given input data I, the teacher generator G t (I) The generated data is represented as I t , student generator G s (I) The generated data is I s , the training images at this time have one-to-one correspondence in structure and texture details; the training strategy is expressed as: in, ‖‖1 is L1 regularization; In step (2), a multi-scale discriminator is used, and the target optimization function of the generative adversarial network is expressed in combination with formula (1) as follows: The discriminator of the student network has the same structure as the discriminator of the teacher network. During training, the weights of the multi-scale discriminator of the teacher network are used to initialize the student discriminator. The weights are not fixed but are adjusted continuously as the compression model is trained. In step (2), feature extraction is performed from the middle layer of the teacher generator. and are the nth intermediate layers to be distilled in the teacher model and the student model, respectively. N represents the total number of layers. The distillation loss is expressed as: Since the number of output channels of the corresponding layers of the teacher model and the student model is different, a feature mapping function F is required before calculation n The output of the intermediate layer of the student model is transformed into a space with the same dimension as the teacher model. The mapping consists of a 1×1 learnable convolutional layer. The distillation process is to optimize and F n To minimize L distill ; In the step (2), when training the student model, a super-net containing multiple candidate channel numbers is first trained, and different channel numbers are arranged and combined to form sub-nets with different structures. The sub-nets and super-net parameters are shared using the once-for-all structure, and the sub-nets obtain the corresponding layer weights from the super-net; the channels of the nth layer excluding the input and output layers of the student network are represented as [C n ], set sub-nets[C n ] The optional channel number range is [16, 24, 32, 48, 64], each layer has 5 channel number structures, and the entire candidate network pool has 5 N Different sub-network structures are randomly selected from the sub-network pool during training, and forward propagation calculation is performed with a certain number of channels, and weights are updated by back propagation according to formula (6); when searching for the optimal network, an evolutionary algorithm is used to screen the sub-network structure. L=L GAN +λ recon L recon +λ distill L distill +λ FM L FM (6) Among them, λ recon ,λ distill ,λ FM These are all hyperparameters; In the step (2), the Fréchet starting distance FID and the dot product operation calculation amount MACs are used as evaluation criteria. Two samples are randomly selected from the candidate pool each time to evaluate FID and MACs, and the candidates with both high values ​​are eliminated. The candidates with lower values ​​are retained as parents, which are fine-tuned to generate offspring and put into the candidate pool until the screening is completed.

2. The method for real-time compensation of color cast and distortion of endoscopic images according to claim 1, characterized in that: In step (3), features are extracted from the non-output layer of the discriminator so that the student model learns the data I generated by the teacher generator t and data generated by the student generator I s Matching of features between; defining sub-discriminator D i The feature extractor of the jth layer is Then the feature matching loss function L FM It is expressed as: Where M is the total number of layers, N j Indicates the number of elements in each layer, D i Not participating in L FM The optimization process only acts as a feature extractor.

3. A real-time compensation device for color cast and distortion of endoscopic images, characterized in that: The device executes the method according to claim 1 and uses a generative adversarial compression network as a standard network, which includes: A pre-training module, which is configured to pre-train a cycle-consistent generative adversarial network with a multi-scale discriminator added to obtain a teacher generator and discriminator; The training module is configured to compress the teacher generator using knowledge distillation and neural structure search, and inherit the multi-scale discriminator to train the student generator; A screening module is configured to evaluate and screen the appropriate compression model structure to obtain the final generator model.

Images