Dense fish school occlusion image restoration method based on diffusion model

Abstract

The invention relates to a dense fish school occlusion image restoration method based on a diffusion model, and belongs to the technical field of image processing. A dense fish shoal occlusion image restoration method based on a diffusion model is characterized in that an improved dense fish shoal occlusion image restoration model is established through the following steps: step 1, image input and coding; step 2, forward diffusion; step 3, condition control; step 4, performing reverse diffusion; and step 5, decoding and image output.

Description

Technical Field

[0001] The invention relates to a method for repairing a dense fish school occluded image based on a diffusion model, and belongs to the technical field of image processing. Background Art

[0002] As aquaculture develops towards intelligence and intensification, it is particularly important to accurately estimate the size of fish. Accurate measurement of fish size can help farmers formulate reasonable fishing quotas, and guide scientific aquaculture management and precise feeding, so as to overcome the eutrophication of the aquaculture water quality environment and protect the aquaculture environment. In traditional aquaculture, the estimation of fish size mainly relies on manual observation and measurement, which is not only inefficient but also highly subjective. With the development of image processing technology, people began to use visual technology to estimate fish size. However, when fish are farmed densely, there are a lot of occlusions, which makes it impossible to directly obtain complete size information from the image. So far, there is no effective method for occlusion recovery of dense fish schools. In order to solve this problem, the present invention proposes a dense fish school occlusion image repair method based on an improved diffusion model. Summary of the invention

[0003] In order to solve the problem of dense fish occlusion restoration, this application provides a dense fish occlusion image restoration method based on a diffusion model. The specific implementation method is as follows: A dense fish school occlusion image restoration method based on a diffusion model, which is special in that an improved dense fish school occlusion image restoration model is established through the following steps: Step 1: Image input and encoding: Input image x0 into the VAE encoder, map image x0 from pixel space to latent space through multiple downsampling, and generate its latent representation z0, z0=E(x0); Step 2, forward diffusion. During the forward diffusion process, noise is gradually added to the potential representation z0 of the image to simulate the transformation process from the potential representation of the clear image to the pure noise potential representation, and finally the pure noise potential representation z is generated. T ; Step 3: Conditional control. The CLIP model is used to extract semantic embedding C(x0) from the input image x0. The embedding contains the shape, color, and texture context information of the target rockfish to be repaired in the image. The concise point prompt p is input by the mouse to indicate the location of the occluded rockfish, which helps the model better focus on the occluded area in the image, thereby improving the repair effect. The concise point prompt p is encoded by the CLIP model to generate a conditional prompt N(p) for conditional control. Step 4, reverse diffusion, the pure noise potential representation z generated by forward diffusion TInput back diffusion; gradually remove noise and repair the area occluded by the spotted sea bream; this process combines conditional control to guide the repair process; generate the potential representation z of the spotted sea bream image 0； Step 5, decoding and image output, input the latent representation z0 of the spotted rockfish image generated by back diffusion into the VAE decoder D(z), decode it and convert it back to pixel space to generate the final repaired spotted rockfish image x1=D(z0).

[0004] Furthermore, in step 2, the noise addition is achieved by the following steps: The addition of noise is related to the time step t, t∈[1,T]. The formula for gradually adding noise is as follows: (1), in, β t is the noise scheduling parameter, which is defined as , β min is the initial noise amount at t=1 in forward diffusion, β max is the maximum noise amount at t=T in forward diffusion; T is the maximum time step, is the noise sampled from a standard Gaussian distribution, ε t ~N(0,1),z t Represents the intermediate state of adding noise. As the time step advances until the maximum time step T, the potential representation z0 of the image is transformed into a pure noise potential representation z T .

[0005] Furthermore, in step 4, during the reverse diffusion process, a denoising UNet is used to gradually denoise the pure noise potential representation; and a switch of a control mechanism is set to control the skip connection operation and the denoising operation.

[0006] Furthermore, the denoising UNet adopts an improved proxy attention mechanism to deal with the problem of repairing the occluded area of ​​densely populated rockfish; the improved proxy attention mechanism uses depthwise separable convolution to focus on the edges and texture areas in image restoration, thereby effectively repairing the edges of the occluded areas in dense fish schools.

[0007] Furthermore, the implementation steps of the improved proxy attention mechanism are as follows: First, the improved proxy attention mechanism calculates the aggregated global information through formula (2): VA : (2), in, VA Represents aggregated global information,A Indicates the proxy token. K is the key, V is the value, σ is the Softmax function, and its formula is as follows: (3), Among them, x i represents the i-th element in the input vector, n represents the dimension of the vector, and e represents the base of the natural logarithm.

[0008] Assign the key K of formula (2) to the conditional prompt N(p), assign the value V to the semantic embedding C(x0), and then calculate VA: (4), Next, the final weighted result O is generated by formula (5): (5), Among them, z t is the intermediate state in the reverse diffusion. A is the proxy token, σ is the Softmax function, VA is the aggregated global information, O is the final weighted result.

[0009] Furthermore, the denoising UNet removes the noise according to formula (6): (6), Formula (6) takes the potential representation z at the current time step t as t As input, combined with the neural network ε θ Predicted noise and noise scheduling parameters β t , calculate the potential representation z of the previous time step t-1 t-1 ; Through gradual updating, the potential representation z from pure noise is finally completed t Transformation to the latent representation z0 of the image.

[0010] Where t is the time step, t∈[1,T], T is the total number of diffusion steps, t represents the current time step, and it gradually decreases from T to 1; z t represents the potential representation of the current time step t, z t-1 represents the potential representation of the previous time step t-1, β t is the noise scheduling parameter, which is defined in the same way as β in formula (1) t , Representation model through neural network ε θ The predicted noise, Ois the weighted result calculated by the improved agent attention, and C(x0) is the semantic embedding extracted from the conditional control.

[0011] Furthermore, it also includes establishing an underwater densely packed rockfish dataset, training models, and image restoration of occluded areas.

[0012] Furthermore, the training model includes an initial repair model, which is trained, weights are updated and parameters are saved to finally obtain an optimized model.

[0013] Further, the processing process of the training model is as follows: During the training process, the model is continuously optimized by calculating the loss function; the loss is calculated after each denoising step, and the model minimizes the prediction noise The error between the actual noise ε is used to update the parameters until convergence; the loss function is in the form of: (7), L represents the loss function, whose value represents the difference between the model's current predicted noise and the actual noise; through back propagation, the model minimizes the loss to adjust the parameters to more accurately predict the noise and generate high-quality images; in, represents the expected value of adding Gaussian noise to z0 at the current time step t; t is the time step, t∈[1,T], the potential representation of the image z0=E(x0), ; The square of the L2 norm ; used to calculate the noise of model predictions The actual noise ε The difference between O is the weighted result calculated by the improved agent attention mechanism, C(x0) represents the semantic embedding extracted from the conditional control, ε is the actual noise sampled from a Gaussian distribution, Representation model through neural network ε θ The predicted noise and θ represent the weights of the model.

[0014] The beneficial effects of the above technical features in this application are: This application proposes a method for repairing images occluded by dense fish schools based on an improved diffusion model. On the basis of establishing a synthetic underwater dense grouper occlusion dataset, the improved diffusion model proxy attention mechanism is used to improve the repair effect in the case of dense grouper occlusion. Through the improved proxy attention mechanism, the model's ability to capture and distinguish key features and semantic information of the image is enhanced, and the model's ability to capture the relationship between distant pixels in the image is enhanced, paying more attention to the edges and texture areas in image repair, thereby improving the quality of the repaired image. This application optimizes the proxy attention mechanism so that it can more effectively handle the repair problem of dense grouper occluded areas. BRIEF DESCRIPTION OF THE DRAWINGS

[0015] Figure 1 A technical roadmap for this application of a dense fish occlusion image restoration method based on an improved diffusion model; Figure 2 This is the structure diagram of the improved agent attention mechanism; Figure 3 This is the original picture of the spotted sea bream; Figure 4 Annotate the dataset diagram; Figure 5 This is a schematic diagram of the underwater densely packed rock snapper dataset; Figure 6 This is the image restoration effect diagram of the method of this application; Figure 7 The repair result diagram of the existing repair method Figure 8 Schematic diagram of the model structure of the method of the present invention. DETAILED DESCRIPTION

[0016] The following embodiments, in conjunction with the accompanying drawings, are only intended to illustrate the technical solutions described in the claims and are not intended to limit the protection scope of the claims.

[0017] The present invention solves the above technical problems through the following technical solutions: Combination Figure 1 A dense fish school occlusion image restoration method based on an improved diffusion model includes establishing an underwater dense grouper dataset, establishing an improved dense fish school occlusion image restoration model, training the model, and restoring the occluded area image.

[0018] Below, taking the spotted sea bream as a specific example, the above scheme is further explained: In order to calculate the accurate number of grouper bream schools, the entire technical route is divided into four parts: a. Establishing a dense underwater grouper bream dataset; b. Building a model; c. Training the model; d. Image repair of occluded areas.

[0019] a. Establish a dataset of densely packed underwater rock snapper: like Figure 3 After collecting underwater images of the grouper using Hikvision underwater cameras, 5,000 images were manually annotated using the ISAT annotation tool. The dataset annotation diagram is shown in the figure. Figure 4 As shown. After the annotation is completed, the Copy-Paste method is used to generate an image with a school of fish occlusion effect: the instance corresponding to the mask part is cut out from one image, and then pasted into another image. The two images are scaled 1:1, and the instance also maintains a 1:1 scale, thereby generating an underwater rockfish image with an occlusion effect. Finally, the dataset is converted into COCO format and divided into a training set and a test set in an 8:2 ratio to complete the establishment of an underwater dense rockfish dataset. Figure 5 Schematic diagram of the underwater densely packed rockfish dataset.

[0020] A method for synthesizing dense fish occlusion datasets is used to generate fish occlusion datasets. This significantly reduces the workload of manual annotation and generates more different occlusion situations, making the position distribution of target objects on the image more consistent with the actual situation.

[0021] b. Build the model: The model structure includes: image input and encoding, forward diffusion, conditional control, reverse diffusion, decoding and image output; the details are as follows: 1. Image input and encoding The image x0 in the dataset is input into the VAE (Variational Autoencoder) encoder. Through multiple downsampling, the image x0 is mapped from the pixel space to the latent space, and its latent representation z0=E(x0) is generated.

[0022] Forward Diffusion In the forward diffusion process, noise is gradually added to the potential representation z0 of the image, simulating the transformation process from the latent representation of the clear image to the pure noise latent representation. The noise added in the process is predefined. Finally, the pure noise latent representation z is generated. T .

[0023] The addition of noise is related to the time step t, t∈[1,T]. The formula for gradually adding noise is as follows: (1), in, β t is the noise scheduling parameter, which is defined as , β min is the initial noise amount at t=1 in forward diffusion, β maxis the maximum amount of noise at t=T in forward diffusion; T is the maximum time step, ε t is the noise sampled from a standard Gaussian distribution, ε t ~N(0,1),z t Represents the intermediate state of adding noise. As the time step advances until the maximum time step T, the potential representation z0 of the image is transformed into a pure noise potential representation z T .

[0024] Condition Control In the process of conditional control, the CLIP model is used to extract semantic embedding C(x0) from the input image x0. This embedding contains the shape, color, and texture context information of the target rockfish to be repaired in the image. The concise point prompt p is input by the mouse to indicate the location of the occluded rockfish, which helps the model better focus on the occluded area in the image, thereby improving the repair effect. The concise point prompt p is encoded by the CLIP model to generate the conditional prompt N(p) for conditional control.

[0025] The semantic embedding C(x0) is combined with the conditional cue N(p) through the following formula (4). This combination guides the attention mechanism to focus on specific areas of the image, thereby providing guidance for back diffusion.

[0026] Back diffusion The pure noise latent representation z generated by forward diffusion T Input back diffusion. It gradually removes noise and inpaints the areas occluded by the spotted rockfish. This process is combined with conditional controls to guide the inpainting process.

[0027] like Figure 8 As shown, in the back diffusion process, the denoising UNet is used to gradually denoise the pure noise potential representation. A switch is set, which is a control mechanism responsible for controlling the jump connection operation and the denoising operation. The jump connection is a mechanism that directly transfers the output of the shallower layer in the denoising UNet to the deeper layer, and the denoising operation refers to the denoising UNet denoising the noisy image. The core part of the denoising UNet is the improved proxy attention mechanism of the present invention.

[0028] Improved proxy attention mechanism: Specific improvements include Figure 2 As shown in the figure, the depthwise separable convolution (DSC) is used to replace the original depthwise convolution (DWC). This improvement significantly reduces the amount of calculation and parameters, and pays more attention to the edges and texture areas in image restoration, thereby effectively repairing the edges of the occluded areas in dense fish schools and improving the restoration speed and quality.

[0029] First, the improved proxy attention mechanism calculates the aggregated global information VA through formula (2): (2), in, VA Represents aggregated global information, A Indicates the proxy token. K is the key, V is the value, σ is the Softmax function, and its formula is as follows: (3), Among them, x i represents the i-th element in the input vector, n represents the dimension of the vector, and e represents the base of the natural logarithm.

[0030] The key in formula (2) K Assign the conditional prompt N(p), value V After assigning the semantic embedding C(x0), VA is calculated: (4), Next, the final weighted result O is generated by formula (5): (5), Among them, z t is the intermediate state in the reverse diffusion. A is the proxy token, σ is the Softmax function, VA is the aggregated global information, and O is the final weighted result.

[0031] Final weighted result O It is used to guide the model to gradually restore the original image. After calculating the final weighted result O Then, we use the denoising UNet to transform the pure noise potential representation z T Start to gradually restore the latent representation z0 of the spotted sea bream image.

[0032] Specifically, the denoising UNet removes the noise according to formula (6): (6), Formula (5) takes the potential representation z at the current time step t as t As input, combined with neural network Predicted noise and noise scheduling parameters β t , calculate the potential representation z of the previous time step t-1 t-1 By gradually updating, the potential representation z from pure noise is finally completed. tTransformation to the latent representation z0 of the image.

[0033] Among them, t is the time step, t∈[1,T], T is the total number of diffusion steps, t represents the current time step, and gradually decreases from T to 1. t represents the potential representation of the current time step t, z t-1 represents the potential representation of the previous time step t-1, β t is the noise scheduling parameter, which is defined in the same way as in formula (1) β t , Representation model through neural network ε θ The predicted noise, O is the weighted result calculated by the improved agent attention, and C(x0) is the semantic embedding extracted from the conditional control.

[0034] By using the above formula, z is continuously updated t-1 , gradually removing the noise until t=1, generating the potential representation z0 of the spotted sea bream image.

[0035] Decoding and image output In the last step of image generation, the latent representation z0 of the spotted rockfish image generated by back diffusion is input into the VAE (Variational Autoencoder) decoder D(z) for decoding and conversion back to pixel space to generate the final restored spotted rockfish image x1=D(z0).

[0036] c. Training model The experiment was implemented on the Pytorch platform, with Python version 3.8 and hardware configuration of three NVIDIA TITAN VGPUs and Intel(R) Xeon(R) Silver 4210 processors. The input image size was 2048×1536, the maximum number of iterations was 10,000, and the batch size was 3. During the training process, the Warmup warm-up strategy was used, with a small initial learning rate. After the model gradually stabilized, it was switched to the preset learning rate, thereby accelerating the model convergence and improving the effect.

[0037] The training module is used to train the dense fish occlusion image restoration model. The specific steps include initial restoration of the model, training, updating weights and saving parameters, and finally obtaining the optimized model.

[0038] During the training process, the model is continuously optimized by calculating the loss function. The loss is calculated after each denoising step. The model minimizes the predicted noise. The error between the actual noise ε is used to update the parameters until convergence. The form of the loss function is: (7), L represents the loss function, whose value represents the difference between the model's current predicted noise and the actual noise. Through back propagation, the model minimizes the loss to adjust the parameters to more accurately predict the noise and generate high-quality images.

[0039] in, represents the expected value of adding Gaussian noise to z0 at the current time step t.

[0040] t is the time step, t∈[1,T], the potential representation of the image z0=E(x0), The square of the L2 norm (i.e., the square of the Euclidean distance) is used to calculate the noise in the model predictions The difference between the actual noise ε. O is the weighted result calculated by the improved proxy attention mechanism, C(x0) represents the semantic embedding extracted from the conditional control, ε is the actual noise sampled from the Gaussian distribution, Representation model through neural network ε θ The predicted noise and θ represent the weights of the model.

[0041] The loss calculation results after each denoising are used to update the model weights θ through back propagation, making the model more accurate in subsequent denoising predictions, thereby generating more realistic images of spotted sea bream.

[0042] d. Image restoration of occluded areas After the model training is completed, the trained model can be used to perform a repair test on an underwater dense rockfish image with occlusions. The model can repair the outline of the occluded area of ​​the rockfish. Based on the repaired image, the size of each rockfish can be calculated.

[0043] After testing, the improved method can accurately restore the complete image outline of the occluded area of ​​the spotted sea bream. The improved method can repair the occluded image outline of the fish more completely, especially in the application of dense underwater fish schools, which lays a technical foundation for the calculation of fish school biomass.

[0044] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principle of the present invention should be included in the protection scope of the present invention.

Claims

1. A dense fish school occlusion image restoration method based on a diffusion model, characterized in that: The improved dense fish occlusion image restoration model is established through the following steps: Step 1: Image input and encoding: Input image x0 into the VAE encoder, map image x0 from pixel space to latent space through multiple downsampling, and generate its latent representation z0, z0=E(x0); Step 2, forward diffusion. During the forward diffusion process, noise is gradually added to the potential representation z0 of the image to simulate the transformation process from the potential representation of the clear image to the pure noise potential representation, and finally the pure noise potential representation z is generated. T ; Step 3: Conditional control. The CLIP model is used to extract semantic embedding C(x0) from the input image x0. The embedding contains the shape, color, and texture context information of the target rockfish to be repaired in the image. The concise point prompt p is input by the mouse to indicate the location of the occluded rockfish, which helps the model better focus on the occluded area in the image, thereby improving the repair effect. The concise point prompt p is encoded by the CLIP model to generate a conditional prompt N(p) for conditional control. Step 4, reverse diffusion, the pure noise potential representation z generated by forward diffusion T Input back diffusion; Gradually remove noise and repair areas occluded by the spotted sea bream; This process combines conditional controls to guide the repair process; Generate the latent representation z of the image of the spotted sea bream 0； Step 5, decoding and image output, input the latent representation z0 of the spotted rockfish image generated by back diffusion into the VAE decoder D(z), decode it and convert it back to pixel space to generate the final repaired spotted rockfish image x1=D(z0).

2. The dense fish school occlusion image restoration method based on the diffusion model according to claim 1 is characterized in that: In step 2, the noise addition is achieved by the following steps: The addition of noise is related to the time step t, t∈[1,T]. The formula for gradually adding noise is as follows: （1）， in, β t is the noise scheduling parameter, which is defined as , β min is the initial noise amount at t=1 in forward diffusion, β max is the maximum noise amount at t=T in forward diffusion; T is the maximum time step, ε t is the noise sampled from a standard Gaussian distribution, ε t ~N(0,1),z t Represents the intermediate state of adding noise. As the time step advances until the maximum time step T, the potential representation z0 of the image is transformed into a pure noise potential representation z T .

3. The dense fish school occlusion image restoration method based on diffusion model according to claim 1 is characterized in that: In the step 4, during the reverse diffusion process, the denoising UNet is used to gradually denoise the pure noise potential representation; and the switch of the control mechanism is set to control the skip connection operation and the denoising operation.

4. The dense fish school occlusion image restoration method based on diffusion model according to claim 3 is characterized in that: The denoising UNet adopts an improved proxy attention mechanism to deal with the problem of repairing the occluded area of ​​densely populated rockfish; the improved proxy attention mechanism uses depthwise separable convolution to focus on the edges and texture areas in image restoration, thereby effectively repairing the edges of the occluded areas in dense fish schools.

5. The dense fish school occlusion image restoration method based on diffusion model according to claim 4 is characterized in that: The implementation steps of the improved proxy attention mechanism are as follows: First, the improved proxy attention mechanism calculates the aggregated global information through formula (2): VA : VA = σ(AK T ）V （2）， in, VA Represents aggregated global information, A Indicates the proxy token. K is the key, V is the value, σ It is the Softmax function, and its formula is as follows: （3）， Among them, x i represents the i-th element in the input vector, n represents the dimension of the vector, and e represents the base of the natural logarithm; Assign the key K of formula (2) to the conditional prompt N(p), and the value V to the semantic embedding C(x0), and then calculate VA: （4）， Next, the final weighted result is generated by formula (5): O : （5）， Among them, z t is the intermediate state in the back diffusion, A is the proxy token, σ is the Softmax function, VA is the aggregated global information, and O is the final weighted result.

6. The dense fish school occlusion image restoration method based on diffusion model according to claim 5 is characterized in that: The denoising UNet removes noise according to formula (6): （6）， Take the potential representation z at the current time step t t As input, combined with neural network Predicted noise and noise scheduling parameters β t , calculate the potential representation z of the previous time step t-1 t-1 ; Through gradual updating, the potential representation z from pure noise is finally completed t Transformation to the latent representation z0 of the image.

7. The dense fish school occlusion image restoration method based on diffusion model according to claim 6 is characterized in that: It also includes establishing a dataset of densely packed underwater rockfish, training models, and image restoration of occluded areas.

8. The dense fish school occlusion image restoration method based on diffusion model according to claim 7 is characterized in that: The training model includes an initial repair model, training, updating weights and saving parameters, and finally obtaining an optimized model.

9. The dense fish school occlusion image restoration method based on diffusion model according to claim 8 is characterized in that: The processing of the training model is as follows: During the training process, the model is continuously optimized by calculating the loss function; the loss is calculated after each denoising step, and the model minimizes the prediction noise The error between the actual noise ε is used to update the parameters until convergence; the loss function is in the form of: （7）， L represents the loss function, whose value represents the difference between the model's current predicted noise and the actual noise; through back propagation, the model minimizes the loss to adjust the parameters to more accurately predict the noise and generate high-quality images; in, represents the expected value of adding Gaussian noise to z0 at the current time step t.

Images