Training method and device for stain detection model and stain detection method and device

Abstract

The present invention discloses a stain detection model training method, device, and stain detection method and device. The stain detection model includes a feature extraction network, a first branch network, and a second branch network. The training method comprises obtaining a training image set; inputting a shell appearance image from the training image set into the stain detection model to obtain predicted stain mask information, stain location information, and stain classification information; obtaining a loss function for the stain detection model based on the predicted stain mask information, stain location information, and stain classification information, as well as label information of the shell appearance image; and adjusting the parameters of the stain detection model based on the loss function of the stain detection model to obtain a trained stain detection model. The trained stain detection model maintains high-resolution information during feature extraction, reduces information loss, and solves the problem of detecting extremely small stains.

Description

Technical Field

[0001] The present invention relates to the field of image processing technology, and in particular to a stain detection model training method and device, and a stain detection method and device. Background Art

[0002] With the continuous advancement of Industry 4.0, computer vision is increasingly being used in industrial inspection. When inspecting product appearance, shell stains are a major cause of substandard product quality. Detecting extremely small stains is particularly challenging due to their small size and varied shapes.

[0003] Methods for detecting extremely small stains are mainly divided into two categories: manual inspection and machine-assisted inspection. Among them, machine-assisted inspection is further divided into traditional visual inspection and deep learning-based visual inspection. Manual inspection is time-consuming and affected by factors such as energy, eyesight, and visual fatigue. At the same time, manual inspection is highly subjective, and each inspection worker has a different understanding of the inspection standards, resulting in inconsistent inspection standards for the same extremely small stain. Relevant machine-assisted inspection has problems such as high computational complexity, slow inspection speed, and sensitivity to spatial displacement and rotation of images. At the same time, when machine-assisted inspection is used for extremely small stains, the stain objects detected vary greatly in shape, size, texture, color, etc., making these methods ineffective for detecting extremely small stains in complex environments. Summary of the Invention

[0004] The present invention aims to at least partially address one of the technical problems in the related art. To this end, one objective of the present invention is to provide a method for training a stain detection model, so that the trained stain detection model maintains high-resolution information during feature extraction, reduces information loss, and solves the problem of detecting extremely small stains.

[0005] The second object of the present invention is to provide a stain detection method.

[0006] The third object of the present invention is to provide a training device for a stain detection model.

[0007] A fourth object of the present invention is to provide a stain detection device.

[0008] A fifth object of the present invention is to provide a computer-readable storage medium.

[0009] A sixth object of the present invention is to provide an electronic device.

[0010] To achieve the above-mentioned purpose, the first embodiment of the present invention proposes a training method for a stain detection model, wherein the stain detection model includes a feature extraction network, a first branch network and a second branch network, wherein the input end of the feature extraction network is used to input an image, the first output end of the feature extraction network is connected to the input end of the first branch network, the second output end of the feature extraction network is connected to the input end of the second branch network, the first branch network is used to output predicted stain mask information, and the second branch network is used to output predicted stain location information and stain classification information. The method comprises: using a training image set to train a stain detection model; The stain detection model is trained multiple times in cycles, wherein for each training cycle, the following operations are performed: a training image set is obtained, wherein the training image set includes a shell appearance image; the shell appearance image in the training image set is input into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information; based on the predicted stain mask information, stain location information and stain classification information, as well as the label information of the shell appearance image, a loss function of the stain detection model is obtained; and according to the loss function of the stain detection model, the parameters of the stain detection model are adjusted to obtain a trained stain detection model.

[0011] According to the stain detection model training method of an embodiment of the present invention, the trained stain detection model always maintains high-resolution information during the feature extraction process, reduces the information loss rate, and solves the task of detecting extremely small stains. The stain detection model adopts a dual-branch structure to make the output results more interpretable.

[0012] To achieve the above-mentioned purpose, the second embodiment of the present invention proposes a stain detection method, which includes: obtaining an image to be tested; inputting the image to be tested into a pre-trained stain detection model to obtain stain mask information, stain position information and stain classification information in the image to be tested, wherein the pre-trained stain detection model is obtained using the training method of the stain detection model proposed in the first embodiment of the present invention.

[0013] According to the stain detection method of an embodiment of the present invention, a trained stain detection model is used to perform extremely small stain detection on the image to be tested. The extremely small stain position information, confidence information, and segmentation mask in the image to be tested can be directly output in the form of an image, making the result more intuitive.

[0014] To achieve the above-mentioned purpose, the third embodiment of the present invention proposes a training device for a stain detection model, wherein the stain detection model includes a feature extraction network, a first branch network and a second branch network, wherein the input end of the feature extraction network is used to input an image, the first output end of the feature extraction network is connected to the input end of the first branch network, the second output end of the feature extraction network is connected to the input end of the second branch network, the first branch network is used to output predicted stain mask information, and the second branch network is used to output predicted stain position information and stain classification information. The device includes: a data acquisition unit, a model training unit, a function calculation unit and a parameter adjustment unit, wherein the data acquisition unit, the model training unit, the function calculation unit and the parameter adjustment unit The entire unit uses a training image set to perform multiple periodic training on the stain detection model, wherein the following operations are performed for each training cycle: a data acquisition unit is used to acquire a training image set, wherein the training image set includes a shell appearance image; a model training unit is used to input the shell appearance image in the training image set into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information; a function calculation unit is used to obtain a loss function of the stain detection model based on the predicted stain mask information, stain location information and stain classification information, as well as the label information of the shell appearance image; a parameter adjustment unit is used to adjust the parameters of the stain detection model according to the loss function of the stain detection model to obtain a trained stain detection model.

[0015] To achieve the above-mentioned purpose, the fourth embodiment of the present invention proposes a stain detection device, which includes: an acquisition module for acquiring an image to be tested; a detection module for inputting the image to be tested into a pre-trained stain detection model to obtain stain mask information, stain position information and stain classification information in the image to be tested, wherein the pre-trained stain detection model is obtained using the training method of the stain detection model proposed in the first embodiment of the present invention.

[0016] To achieve the above-mentioned objectives, the fifth embodiment of the present invention proposes a computer-readable storage medium on which a computer program is stored. When the computer program is executed by a processor, it implements the training method of the stain detection model proposed in the first embodiment of the present invention, or the stain detection method proposed in the second embodiment of the present invention.

[0017] To achieve the above-mentioned objectives, the sixth embodiment of the present invention proposes an electronic device, including a memory and a processor, wherein a computer program is stored on the memory. When the computer program is executed by the processor, the training method of the stain detection model proposed in the first embodiment of the present invention, or the stain detection method proposed in the second embodiment of the present invention is implemented.

[0018] Additional aspects and advantages of the present invention will be set forth in part in the description which follows and, in part, will be obvious from the description which follows, or may be learned through practice of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 A schematic diagram of a very small stain according to an embodiment of the present invention;

[0020] Figure 2 A schematic diagram of a stain detection model according to an embodiment of the present invention;

[0021] Figure 3 This is a flowchart of a method for training a stain detection model according to an embodiment of the present invention;

[0022] Figure 4 This is a flowchart of a method for training a stain detection model according to a specific embodiment of the present invention;

[0023] Figure 5 A schematic diagram of a convolutional neural network according to an embodiment of the present invention;

[0024] Figure 6 is a schematic diagram of Stage 2 according to an embodiment of the present invention;

[0025] Figure 7 Schematic diagram of a Stage 4-1 convolution module according to an embodiment of the present invention;

[0026] Figure 8 Schematic diagram of a Stage 4-3 convolution module according to an embodiment of the present invention;

[0027] Figure 9 Schematic diagram of a BasicBlock convolution block according to an embodiment of the present invention;

[0028] Figure 10 Schematic diagram of a BasicBlock_SE convolution block according to an embodiment of the present invention;

[0029] Figure 11 A schematic diagram of an ECA module according to an embodiment of the present invention;

[0030] Figure 12 Schematic diagram of an SE module according to an embodiment of the present invention;

[0031] FIG13( a ) is a schematic diagram showing the before and after effects of stain detection using a trained stain detection model according to a specific embodiment of the present invention;

[0032] FIG13( b ) is a schematic diagram showing the before and after effects of stain detection using a trained stain detection model according to another specific embodiment of the present invention;

[0033] Figure 14 This is a flow chart of a stain detection method according to an embodiment of the present invention;

[0034] Figure 15 A schematic diagram of a training device for a stain detection model according to an embodiment of the present invention;

[0035] Figure 16 FIG. 1 is a schematic diagram of a stain detection device according to an embodiment of the present invention. DETAILED DESCRIPTION

[0036] The following describes embodiments of the present invention in detail, examples of which are shown in the accompanying drawings, wherein the same or similar reference numerals throughout represent the same or similar elements or elements having the same or similar functions. The embodiments described below with reference to the accompanying drawings are exemplary and are intended to be used to explain the present invention, and are not to be construed as limiting the present invention.

[0037] The accuracy of industrial cameras currently used in the industrial field is 0.01mm / pixel. When the front size of the shell is 98*73mm, the size of the smallest stain is between 0.1*0.1-0.4*0.6mm. The digital image resolution is 1200*1000. Therefore, the pixels corresponding to the smallest stain are approximately between 1*1-4*4 pixels. Figure 1 It should be noted that Figure 1 The small and medium dashed boxes represent very small stains in real-world scenarios. For ease of illustration, the large dashed box shows the stain magnified five times. Related detection methods or convolutional neural networks cannot effectively extract the characteristics of these stains.

[0038] The present invention provides a stain detection model training method, device, and stain detection method and device to solve the above problems. The stain detection model training method, device, and stain detection method and device of the embodiment of the present invention are described in detail below in conjunction with instructions 1-16 and specific implementation methods.

[0039] The stain detection model in the embodiment of the present invention is as follows: Figure 2 As shown, the stain detection model may include a feature extraction network, a first branch network and a second branch network, wherein the input end of the feature extraction network is used to input an image, the first output end of the feature extraction network is connected to the input end of the first branch network, the second output end of the feature extraction network is connected to the input end of the second branch network, the first branch network is used to output predicted stain mask information, and the second branch network is used to output predicted stain position information and stain classification information.

[0040] To improve the interpretability of the stain detection model, the stain detection model in this embodiment of the present invention adopts a dual-branch structure, consisting of a first branch network for segmentation and a second branch network for category recognition. This dual-branch design cleverly decouples the classification and segmentation feature layers, with different branches responsible for specialized tasks. This information is then directly displayed as an image on the terminal, helping inspectors quickly locate and screen the smallest stain shells.

[0041] In an embodiment of the present invention, a feature extraction network extracts features from an input image. The resulting common feature map is fed into a first branch network and a second branch network, respectively. The first branch network outputs predicted stain mask information for the input image based on the common feature map, while the second branch network outputs predicted stain location information and stain classification information for the input image based on the common feature map. The stain mask, stain location, and stain classification information are displayed as images, making the output results more intuitive.

[0042] In an embodiment of the present invention, when building a stain detection model, the language that can be used is Python, and the deep learning framework is Pytorch.

[0043] In an embodiment of the present invention, the training method of the stain detection model includes: using the training image set to perform multiple periodic training on the stain detection model, wherein for each training cycle, the following operations may be performed, see Figure 3 :

[0044] S1, obtaining a training image set, wherein the training image set includes an image of the shell appearance.

[0045] Specifically, an appearance picture of the housing captured by an industrial camera is obtained. When the appearance picture of the housing is captured by the industrial camera, the housing can be placed in a viewfinder, and the industrial camera in the viewfinder illuminates and photographs the front and side of the housing.

[0046] In an embodiment of the present invention, acquiring the training image set may include: collecting an appearance image of the shell, and preprocessing the appearance image of the shell.

[0047] To accelerate the training and inference of the stain detection model, the data size of the housing appearance image captured by the industrial camera can be controlled within a reasonable range. Otherwise, the data feature extraction of the stain detection model will be time-consuming, which is not conducive to efficient production. For example, the housing appearance image can be scaled from 3000*3000 to 1200*1000.

[0048] To enhance the robustness and generalization of the model, the shell appearance image was digitized, flipped at random angles and mirrored, and randomly edge-filled. Taking into account shooting conditions and lighting issues, the resulting image was randomly inverted and random noise was added.

[0049] In an embodiment of the present invention, the obtained housing appearance images can be divided into a training image set and a test image set in a ratio of 8:2. During the training phase, the housing appearance images in the training image set are used to train the stain detection model. After training, the trained stain detection model is evaluated using the housing appearance images in the test image set.

[0050] S2: Input the shell appearance image in the training image set into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information.

[0051] Specifically, the feature extraction network is used to extract features from the input shell appearance image and output a common feature map of the shell appearance image. This common feature map is input into the first branch network to obtain stain mask information for the input shell appearance image. This common feature map is input into the second branch network to obtain stain location information and stain classification information for the input shell appearance image.

[0052] In an embodiment of the present invention, Figure 4 As shown, the feature extraction network may include a convolutional neural network and a candidate box extraction network. The shell appearance image in the training image set is input into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information, which may include:

[0053] S21, inputting the shell appearance image into a convolutional neural network to obtain a common feature map of the shell appearance image;

[0054] S22, using the candidate box extraction network to slide on the common feature map to obtain a common feature map with stain candidate areas;

[0055] S23, inputting the common feature map with the stain candidate area into the first branch network to obtain predicted stain mask information;

[0056] S24: Input the common feature map with the stain candidate area into the second branch network to obtain predicted stain location information and predicted stain classification information.

[0057] Because the classification and location of very small stains have highly overlapping feature information, in order to ensure that the feature extraction network maintains high-resolution information during the extraction process, reduce information loss, and solve the very small stain detection task, in some embodiments, the shell appearance image is input into a convolutional neural network to obtain a common feature map. The candidate box extraction network is then slid across the common feature map to obtain candidate areas where stains may appear. These candidate areas are then sent to the first branch network and the second branch network for branch decoding, respectively, to obtain a stain segmentation mask, stain location information, and stain classification information.

[0058] In some embodiments, the first branch network may be a mask branch, and the second branch network may be a fully connected layer.

[0059] S3, based on the predicted stain mask information, stain location information and stain classification information, as well as the label information of the shell appearance image, obtain the loss function of the stain detection model.

[0060] S4, adjusting the parameters of the stain detection model according to the loss function of the stain detection model to obtain a trained stain detection model.

[0061] During training, a loss function is constructed for the stain detection model. The loss value is calculated using the measured stain mask, location, and classification information, as well as the label information of the shell appearance image. Based on this loss value, the parameters of the stain detection model are adjusted. During training, the loss function is used to calculate the loss value for each shell appearance image. The loss value is minimized during training. When the loss value is minimized, the trained stain detection model is obtained.

[0062] The trained stain detection model is evaluated using the shell appearance images from the test image set. Evaluation metrics include average box precision, classification accuracy, and average pixel segmentation accuracy. Box precision represents the model's ability to find small stains, classification accuracy represents the model's ability to classify background and very small stains, and pixel segmentation accuracy represents the model's ability to distinguish individual pixels. After the evaluation is complete, the resulting training weights are saved, and the training phase ends.

[0063] It should be noted that the training method of the stain detection model in the embodiment of the present invention includes two stages, namely a training stage and an inference stage.

[0064] In the embodiment of the present invention, the inference stage directly inputs the collected shell appearance image into the trained model weights without any flipping, filling, inversion, or noise processing to obtain the output result.

[0065] In some embodiments, the constructed stain detection model is trained on a Linux Ubuntu 18.6 platform and a NVIDIA Tesla V100*2 graphics card. The training parameters are as follows: Lr = 0.0025, momentum = 0.9, weight_decay = 0.0001, optimizer = SGD, max_epoch = 60, and batch_size = 2.

[0066] In some embodiments, the convolutional neural network may adopt an improved HRNetv3 convolutional neural network.

[0067] In an embodiment of the present invention, Figure 5 As shown, the convolutional neural network may include a first convolution module, a second convolution module (Stage 2), a third convolution module (Stage 3-1, Stage 3-2, Stage 3-3 and Stage 3-4) and a fourth convolution module (Stage 4-1, Stage 4-2 and Stage 4-4) connected in sequence. 4-3), the input end of the first convolution module is used to input the image, the first output end of the first convolution module is connected to the first input end of the second convolution module, and the second output end of the first convolution module is connected to the second input end of the second convolution module: the first output end of the second convolution module is connected to the first input end of the third convolution module, the second output end of the second convolution module is connected to the second input end of the third convolution module, and the third output end of the second convolution module is connected to the third input end of the third convolution module; the first output end of the third convolution module is connected to the first input end of the fourth convolution module, the second output end of the third convolution module is connected to the second input end of the fourth convolution module, the third output end of the third convolution module is connected to the third input end of the fourth convolution module, and the fourth output end of the third convolution module is connected to the fourth input end of the fourth convolution module; the output end of the fourth convolution module outputs a common feature map.

[0068] In some embodiments, the convolutional neural network is divided into four stages, each with a similar structure. Their purpose is to ensure that the shell appearance image is always able to incorporate high-resolution feature information during feature extraction, thereby preventing information loss during downsampling. Compared to the original network (HRNetv2 convolutional neural network), the convolutional neural network of the embodiment of the present invention adds one stage, increasing the depth and width of the original network (HRNetv2 convolutional neural network).

[0069] in, Figure 6The figure shows the feature extraction process of the second convolutional module (Stage 2). The original branches are sized 64*64 (representing high resolution) and 32*32 (representing low resolution). Both branches pass through five BasicBlock convolution blocks, and their feature sizes remain unchanged. At this time, the high-resolution feature layer is convolutionally downsampled to 32*32 size and then fused with the 32*32 feature map in the original branch. The low-resolution feature layer is upsampled to 64*64 size and fused with the 64*64 feature map in the original branch to always retain high-resolution information.

[0070] It should be noted that the third convolution module (Stage 3-1, Stage 3-2, Stage 3-3 and Stage 3-4) is similar to the second convolution module (Stage 2), but has a different width.

[0071] In one embodiment of the present invention, Figure 7 and Figure 8 As shown, the fourth convolution module may include three convolution components connected in sequence, which are respectively recorded as the first convolution component (Stage 4-1), the second convolution component (Stage 4-2) and the third convolution component (Stage 4-3); wherein, the first convolution component (Stage 4-1) and the second convolution component (Stage 4-2) may include four convolution branches, which are respectively recorded as the first convolution branch, the second convolution branch, the third convolution branch and the fourth convolution branch, and the first convolution branch includes four BasicBlock_SE convolution blocks connected in sequence (see Figure 7 ), the second convolution branch includes four BasicBlock_SE convolution blocks and a 2x upsampling block connected in sequence (see Figure 7 ), the third convolution branch includes four BasicBlock_SE convolution blocks and a 4x upsampling block connected in sequence (see Figure 7 ), the fourth convolution branch includes four BasicBlock_SE convolution blocks and an 8-fold upsampling block connected in sequence (see Figure 7 ); The third convolution component (Stage 4-3) includes four convolution branches, a fusion layer, and an activation layer, which are respectively recorded as the fifth convolution branch, the sixth convolution branch, the seventh convolution branch, and the eighth convolution branch. The fifth convolution branch includes four sequentially connected BasicBlock_SE convolution blocks and an efficient channel attention module. The sixth convolution branch, the seventh convolution branch, and the eighth convolution branch all include four sequentially connected BasicBlock_SE convolution blocks, a convolution layer, a batch normalization layer, an upsampling layer, and an efficient channel attention module (see Figure 8 ).

[0072] In some embodiments, the structure of the first convolutional component is as follows Figure 7 As shown, the original branches are 64*64 (representing high resolution), 32*32 (representing medium resolution), 16*16 (representing low resolution), and 8*8 (representing ultra-low resolution). All four branches pass through four BasicBlock_SE convolution blocks. The structure is as follows Figure 9 As shown in the figure, the feature size of the first branch (64*64) remains unchanged after convolution. The other three branches are sampled by four BasicBlock_SE convolution blocks and then upsampled by 2, 4, and 8 times, respectively, to obtain new feature maps. At this time, the resolution of the four branches is 64*64. The feature maps of the three upsampled branches are fused with the 64*64 feature map of the original branch.

[0073] In an embodiment of the present invention, the second convolution component of the fourth convolution module has the same structure as the first convolution component, and the result of the second convolution component is input to the third convolution component to obtain the final common feature map. The final common feature map can always retain high-resolution information.

[0074] In an embodiment of the present invention, Figure 9 As shown, the BasicBlock convolution block may include a 2D convolution layer, a batch normalization layer, an activation layer, a 2D convolution layer, a batch normalization layer, a fusion layer and an activation layer connected in sequence.

[0075] In an embodiment of the present invention, Figure 10 As shown, the BasicBlock_SE convolution block may include a 2D convolution layer, a batch normalization layer, an activation layer, a 2D convolution layer, a compressed excitation attention module, a fusion layer and an activation layer connected in sequence.

[0076] In an embodiment of the present invention, in order to enhance the connection and information connectivity between feature channels, strengthen the information signals of important channels in the feature map, and weaken the signals of unimportant channels in the feature map, a local cross-channel interaction strategy without dimensionality reduction is proposed. An efficient channel attention module (ECA) and a squeeze excitation attention module (SE) are added to the network. Appropriate cross-channel interaction can significantly reduce the complexity of the model while maintaining performance, effectively avoiding the impact of dimensionality reduction on the channel attention learning effect.

[0077] In an embodiment of the present invention, Figure 11As shown in the figure, the ECA module performs a global average pooling operation on the input feature map and a 1D convolution operation with an adaptive kernel size of k. The weight w of each channel is obtained by the Sigmoid activation function. The formula is as follows:

[0078] C=φ(k)=2 (γ*k-b)

[0079] Where C is known and is the number of channels in the feature layer of the input module. In the embodiment of the present invention, γ is set to 2 and b is set to 1. Then, k can be calculated by the above formula. The k convolution size is:

[0080]

[0081] If k is calculated to be a decimal or an even number, then take the odd number closest to k. Multiply the weight by the corresponding element of the original input feature map to obtain the final output feature map. It should be noted that Figure 11 in Represents element-wise product.

[0082] In an embodiment of the present invention, Figure 12 As shown in Figure 1, the SE module operation steps are divided into three steps: compression operation, excitation operation and scale operation.

[0083] Among them, the feature size of the compression operation input is: W*H*C. After a global average pooling, the feature map is compressed to a 1*1*C vector after the compression operation. The excitation operation consists of two fully connected layers. The first fully connected layer has C*SE-Ratio neurons, the input is 1×1×C, and the output is 1×1×C×SE-Radio. The second fully connected layer has C neurons, the input is 1×1×C×SERadio, and the output is 1×1×C, where SE-Ratio is a scaling parameter. The purpose of this parameter is to reduce the number of channels and thus reduce the amount of calculation. It is a positive integer, and 16 is taken in this invention. After the scale operation obtains the 1×1×C vector, the original feature map can be scaled. It is very simple, that is, the channel weights are multiplied. The original feature vector is W×H×C. The weight values ​​of each channel calculated by the SE module are multiplied by the two-dimensional matrix of the corresponding channel of the original feature map, and the resulting output is obtained.

[0084] The stain detection model trained in the embodiment of the present invention can detect and segment different types of very small stains and output the detection results without adding a new first branch network. When the types of very small stains increase, only retraining is required to complete the detection of new stains without modifying a large amount of code. The stain detection model can identify very small stain shells. The stain detection model can simultaneously output the mask information, location information, and confidence information of very small stains. As shown in Figures 13(a) and 13(b), the stain detection model can directly draw the mask of very small stains on the original image, which helps the inspection personnel to quickly locate and screen very small stains. Among them, the left images in Figures 13(a) and 13(b) are the original images input to the model (the stains are too small to be seen by the naked eye), the middle images in Figures 13(a) and 13(b) are the output images of the trained stain detection model with segmentation masks and location boxes, and the right images in Figures 13(a) and 13(b) are schematic diagrams of a stain in the original image after it is enlarged for ease of display.

[0085] The stain detection model training method of the present invention solves the problem of detecting extremely small stains. The super-resolution representation learning employed by the stain detection model maintains high-resolution information throughout the extraction process, reducing information loss and enabling the detection of extremely small stains. The stain detection model employs a dual-branch structure, making the model results more interpretable. The stain detection model outputs information about the location of extremely small stains, confidence levels, and segmentation masks, making the results more intuitive. When a very small stain is detected, a dark mask is added to the image at the stain location, enabling rapid location identification and facilitating practical applications.

[0086] The invention provides a stain detection method.

[0087] Figure 14 FIG. 1 is a flow chart of a stain detection method according to an embodiment of the present invention. Figure 14 As shown, the stain detection method may include:

[0088] S31, obtaining an image to be tested;

[0089] S32, inputting the image to be tested into a pre-trained stain detection model to obtain stain mask information, stain location information and stain classification information in the image to be tested, wherein the pre-trained stain detection model is obtained using the training method of the stain detection model as described above.

[0090] Specifically, the system obtains an image of the housing to be inspected and feeds it into a pre-trained stain detection model. This model then generates stain mask information, stain location information, and stain classification information. When very small stains are detected, a dark mask is added to the image at the stain location, enabling rapid location tracking.

[0091] It should be noted that, for other specific implementations of the stain detection method provided in the embodiment of the present invention, reference may be made to other specific implementations of the stain detection model training method in the above embodiment of the present invention.

[0092] The stain detection method of the embodiment of the present invention uses a trained stain detection model to detect extremely small stains in the image to be tested. The extremely small stain position information, confidence information, and segmentation mask in the image to be tested can be directly output in the form of an image, making the result more intuitive.

[0093] The present invention provides a training device for a stain detection model.

[0094] In an embodiment of the present invention, the stain detection model may include a feature extraction network, a first branch network and a second branch network, wherein the input end of the feature extraction network is used to input an image, the first output end of the feature extraction network is connected to the input end of the first branch network, the second output end of the feature extraction network is connected to the input end of the second branch network, the first branch network is used to output predicted stain mask information, and the second branch network is used to output predicted stain location information and stain classification information.

[0095] Figure 15 FIG. 1 is a schematic diagram of a training device for a stain detection model according to an embodiment of the present invention. Figure 15 As shown, the training device 100 for the stain detection model includes: a data acquisition unit 10, a model training unit 20, a function calculation unit 30 and a parameter adjustment unit 40.

[0096] Among them, the data acquisition unit 10, the model training unit 20, the function calculation unit 30 and the parameter adjustment unit 40 use the training image set to perform multiple periodic training on the stain detection model, wherein for each training cycle, the following operations are performed: the data acquisition unit 10 is used to acquire the training image set, wherein the training image set includes a shell appearance image; the model training unit 20 is used to input the shell appearance image in the training image set into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information; the function calculation unit 30 is used to obtain the loss function of the stain detection model based on the predicted stain mask information, stain location information and stain classification information, as well as the label information of the shell appearance image; the parameter adjustment unit 40 is used to adjust the parameters of the stain detection model according to the loss function of the stain detection model to obtain a trained stain detection model.

[0097] In one embodiment of the present invention, the data acquisition unit is specifically configured to collect a shell appearance image and pre-process the shell appearance image.

[0098] In one embodiment of the present invention, the feature extraction network includes a convolutional neural network and an extraction candidate box network, and a model training unit, which is specifically used to input the shell appearance image into the convolutional neural network to obtain a common feature map of the shell appearance image; use the extraction candidate box network to slide on the common feature map to obtain a common feature map with a stain candidate area; input the common feature map with the stain candidate area into the first branch network to obtain predicted stain mask information; input the common feature map with the stain candidate area into the second branch network to obtain predicted stain position information and predicted stain classification information.

[0099] In one embodiment of the present invention, the convolutional neural network includes a first convolution module, a second convolution module (Stage 2), a third convolution module (Stage 3-1, Stage 3-2, Stage 3-3 and Stage 3-4) and a fourth convolution module (Stage 4-1, Stage 4-2 and Stage 4-4) connected in sequence. 4-3), the input end of the first convolution module is used to input the image, the first output end of the first convolution module is connected to the first input end of the second convolution module, and the second output end of the first convolution module is connected to the second input end of the second convolution module: the first output end of the second convolution module is connected to the first input end of the third convolution module, the second output end of the second convolution module is connected to the second input end of the third convolution module, and the third output end of the second convolution module is connected to the third input end of the third convolution module; the first output end of the third convolution module is connected to the first input end of the fourth convolution module, the second output end of the third convolution module is connected to the second input end of the fourth convolution module, the third output end of the third convolution module is connected to the third input end of the fourth convolution module, and the fourth output end of the third convolution module is connected to the fourth input end of the fourth convolution module; the output end of the fourth convolution module outputs a common feature map.

[0100] In some embodiments, the convolutional neural network is divided into four stages, each with a similar structure. Their purpose is to ensure that the shell appearance image is always able to incorporate high-resolution feature information during feature extraction, thereby preventing information loss during downsampling. Compared to the original network (HRNetv2 convolutional neural network), the convolutional neural network of the embodiment of the present invention adds one stage, increasing the depth and width of the original network (HRNetv2 convolutional neural network).

[0101] in, Figure 6The figure shows the feature extraction process of the second convolutional module (Stage 2). The original branches are sized 64*64 (representing high resolution) and 32*32 (representing low resolution). Both branches pass through five BasicBlock convolution blocks, and their feature sizes remain unchanged. At this time, the high-resolution feature layer is convolutionally downsampled to 32*32 size and then fused with the 32*32 feature map in the original branch. The low-resolution feature layer is upsampled to 64*64 size and fused with the 64*64 feature map in the original branch to always retain high-resolution information.

[0102] It should be noted that the third convolution module (Stage 3-1, Stage 3-2, Stage 3-3 and Stage 3-4) is similar to the second convolution module (Stage 2), but has a different width.

[0103] In one embodiment of the present invention, Figure 7 and Figure 8 As shown, the fourth convolution module may include three convolution components connected in sequence, which are respectively recorded as the first convolution component (Stage 4-1), the second convolution component (Stage 4-2) and the third convolution component (Stage 4-3); wherein, the first convolution component (Stage 4-1) and the second convolution component (Stage 4-2) may include four convolution branches, which are respectively recorded as the first convolution branch, the second convolution branch, the third convolution branch and the fourth convolution branch, and the first convolution branch includes four BasicBlock_SE convolution blocks connected in sequence (see Figure 7 ), the second convolution branch includes four BasicBlock_SE convolution blocks and a 2x upsampling block connected in sequence (see Figure 7 ), the third convolution branch includes four BasicBlock_SE convolution blocks and a 4x upsampling block connected in sequence (see Figure 7 ), the fourth convolution branch includes four BasicBlock_SE convolution blocks and an 8-fold upsampling block connected in sequence (see Figure 7 ); The third convolution component (Stage 4-3) includes four convolution branches, a fusion layer, and an activation layer, which are respectively recorded as the fifth convolution branch, the sixth convolution branch, the seventh convolution branch, and the eighth convolution branch. The fifth convolution branch includes four sequentially connected BasicBlock_SE convolution blocks and an efficient channel attention module. The sixth convolution branch, the seventh convolution branch, and the eighth convolution branch all include four sequentially connected BasicBlock_SE convolution blocks, a convolution layer, a batch normalization layer, an upsampling layer, and an efficient channel attention module (see Figure 8 ).

[0104] In one embodiment of the present invention, the BasicBlock convolution block may include a 2D convolution layer, a batch normalization layer, an activation layer, a 2D convolution layer, a batch normalization layer, a fusion layer and an activation layer connected in sequence.

[0105] In one embodiment of the present invention, the BasicBlock_SE convolution block may include a 2D convolution layer, a batch normalization layer, an activation layer, a 2D convolution layer, a compressed excitation attention module, a fusion layer and an activation layer connected in sequence.

[0106] It should be noted that other specific implementations of the stain detection model training device provided in the embodiment of the present invention can refer to other specific implementations of the stain detection model training method in the above embodiment of the present invention.

[0107] The stain detection model training device of the present invention solves the problem of detecting extremely small stains. The super-resolution representation learning employed by the stain detection model can consistently maintain high-resolution information during the extraction process, reducing information loss and addressing the task of detecting extremely small stains. The stain detection model employs a dual-branch structure, making the model results more interpretable. The stain detection model can output extremely small stain location information, confidence information, and a segmentation mask, making the results more intuitive. When an extremely small stain is detected, a dark mask is added to the image at the stain location, enabling rapid location of the stain and facilitating practical applications.

[0108] The invention provides a stain detection device.

[0109] Figure 16 FIG. 1 is a schematic diagram of a stain detection device according to an embodiment of the present invention. Figure 16 The stain detection device 200 may include: an acquisition module 50 and a detection module 60 .

[0110] Among them, the acquisition module 50 is used to acquire the image to be tested; the detection module 60 is used to input the image to be tested into a pre-trained stain detection model to obtain stain mask information, stain location information and stain classification information in the image to be tested, wherein the pre-trained stain detection model is obtained using the training method of the stain detection model as described above.

[0111] It should be noted that, for other specific implementations of the stain detection device provided in the embodiment of the present invention, reference may be made to other specific implementations of the stain detection model training method in the above embodiment of the present invention.

[0112] The stain detection device of the embodiment of the present invention uses a trained stain detection model to detect extremely small stains on the image to be tested, and can directly output the extremely small stain position information, confidence information, and segmentation mask in the image to be tested in the form of an image, making the result more intuitive.

[0113] The present invention provides a computer-readable storage medium.

[0114] In this embodiment, a computer program is stored on a computer-readable storage medium. When the computer program is executed by a processor, the training method of the stain detection model as described above is implemented.

[0115] The present invention provides an electronic device.

[0116] In this embodiment, the electronic device includes a memory and a processor. The memory stores a computer program. When the computer program is executed by the processor, the training method of the stain detection model described above, or the stain detection method described above, is implemented.

[0117] The computer-readable storage medium and electronic device in the embodiments of the present invention utilize the aforementioned stain detection model training method to obtain a trained stain detection model. Using the aforementioned stain detection method, the trained stain detection model can be used to perform minimal stain detection on a test image. The location information, confidence level information, and segmentation mask of the minimal stain in the test image can be directly output as an image, making the results more intuitive.

[0118] It should be noted that the logic and / or steps represented in the flowcharts or otherwise described herein, for example, can be considered as a sequenced list of executable instructions for implementing the logical functions, and can be embodied in any computer-readable medium for use by an instruction execution system, apparatus, or device (such as a computer-based system, a system including a processor, or other system that can fetch and execute instructions from an instruction execution system, apparatus, or device), or in conjunction with such instruction execution system, apparatus, or device. For the purposes of this specification, a "computer-readable medium" can be any device that can contain, store, communicate, propagate, or transmit a program for use by an instruction execution system, apparatus, or device, or in conjunction with such instruction execution system, apparatus, or device. More specific examples (non-exhaustive list) of computer-readable media include the following: an electrical connection portion having one or more wires (electronic device), a portable computer disk cartridge (magnetic device), random access memory (RAM), read-only memory (ROM), erasable and programmable read-only memory (EPROM or flash memory), fiber optic devices, and portable compact disk read-only memory (CDROM). Furthermore, the computer-readable medium may even be paper or other suitable medium on which the program is printed, since the program may be obtained electronically, for example, by optically scanning the paper or other medium and then editing, interpreting or processing it in another suitable manner if necessary, and then storing it in a computer memory.

[0119] It should be understood that various parts of the present invention can be implemented using hardware, software, firmware, or a combination thereof. In the above-described embodiments, multiple steps or methods can be implemented using software or firmware stored in a memory and executed by a suitable instruction execution system. For example, if implemented using hardware, as in another embodiment, any one of the following technologies known in the art or a combination thereof can be used: a discrete logic circuit having a logic gate circuit for implementing a logic function on a data signal, an application-specific integrated circuit having a suitable combination of logic gate circuits, a programmable gate array (PGA), a field programmable gate array (FPGA), etc.

[0120] Throughout this specification, reference to terms such as "one embodiment," "some embodiments," "examples," "specific examples," or "some examples" means that a specific feature, structure, material, or characteristic described in conjunction with that embodiment or example is included in at least one embodiment or example of the present invention. In this specification, schematic representations of the above terms do not necessarily refer to the same embodiment or example. Furthermore, the specific features, structures, materials, or characteristics described may be combined in any suitable manner in any one or more embodiments or examples.

[0121] In the description of the present invention, it should be understood that the terms "center", "longitudinal", "lateral", "length", "width", "thickness", "up", "down", "front", "back", "left", "right", "vertical", "horizontal", "top", "bottom", "inside", "outside", "clockwise", "counterclockwise", "axial", "radial", "circumferential" and the like to indicate orientations or positional relationships based on the orientations or positional relationships shown in the accompanying drawings, and are only for the convenience of describing the present invention and simplifying the description, rather than indicating or implying that the device or element referred to must have a specific orientation, be constructed and operated in a specific orientation, and therefore should not be understood as limiting the present invention.

[0122] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of the technical features being referred to. Thus, a feature defined as "first" or "second" may explicitly or implicitly include at least one such feature. In the description of the present invention, "plurality" means at least two, such as two, three, etc., unless otherwise specifically defined.

[0123] In the present invention, unless otherwise specified or limited, the terms "installed," "connected," "connect," "fixed," etc. should be understood in a broad sense. For example, they can refer to fixed connection, detachable connection, or integration; mechanical connection, electrical connection; direct connection, or indirect connection through an intermediate medium; internal communication between two components, or interaction between two components, unless otherwise specified. Those skilled in the art will understand the specific meanings of the above terms in the present invention based on specific circumstances.

[0124] In the present invention, unless otherwise expressly specified or limited, when a first feature is "above" or "below" a second feature, it may mean that the first and second features are in direct contact, or that the first and second features are in indirect contact through an intermediary. Furthermore, when a first feature is "above," "above," or "above" a second feature, it may mean that the first feature is directly above or diagonally above the second feature, or simply means that the first feature is at a higher level than the second feature. When a first feature is "below," "below," or "below" a second feature, it may mean that the first feature is directly below or diagonally below the second feature, or simply means that the first feature is at a lower level than the second feature.

[0125] Although the embodiments of the present invention have been shown and described above, it will be understood that the above embodiments are illustrative and are not to be construed as limitations on the present invention. A person skilled in the art may change, modify, replace and modify the above embodiments within the scope of the present invention.

Claims

1. A stain detection model training method, characterized in that: The stain detection model includes a feature extraction network, a first branch network, and a second branch network, wherein an input end of the feature extraction network is used to input an image, a first output end of the feature extraction network is connected to the input end of the first branch network, and a second output end of the feature extraction network is connected to the input end of the second branch network, the first branch network is used to output predicted stain mask information, and the second branch network is used to output predicted stain location information and stain classification information. The method includes: The stain detection model is trained multiple times using the training image set. For each training cycle, the following operations are performed: Acquire a training image set, wherein the training image set includes an image of the shell appearance; Inputting the shell appearance image in the training image set into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information; Obtaining a loss function of the stain detection model based on the predicted stain mask information, stain location information, and stain classification information, as well as label information of the shell appearance image; Adjusting the parameters of the stain detection model according to the loss function of the stain detection model to obtain a trained stain detection model; The feature extraction network includes a convolutional neural network and a candidate box extraction network. The shell appearance image in the training image set is input into the stain detection model to obtain predicted stain mask information, stain location information and stain classification information, including: Inputting the shell appearance image into the convolutional neural network to obtain a common feature map of the shell appearance image; Utilize the candidate frame extraction network to slide on the common feature map to obtain a common feature map having a stain candidate area; Inputting the common feature map having the stain candidate area into the first branch network to obtain the predicted stain mask information; The common feature map having the stain candidate area is input into the second branch network to obtain the predicted stain position information and the predicted stain classification information.

2. The stain detection model training method according to claim 1, characterized in that: Obtaining a training image set includes: The shell appearance picture is collected and preprocessed.

3. The stain detection model training method according to claim 1, characterized in that: The convolutional neural network includes a first convolution module, a second convolution module, a third convolution module and a fourth convolution module connected in sequence. The input end of the first convolution module is used to input an image, the first output end of the first convolution module is connected to the first input end of the second convolution module, and the second output end of the first convolution module is connected to the second input end of the second convolution module: The first output terminal of the second convolution module is connected to the first input terminal of the third convolution module, the second output terminal of the second convolution module is connected to the second input terminal of the third convolution module, and the third output terminal of the second convolution module is connected to the third input terminal of the third convolution module; The first output of the third convolution module is connected to the first input of the fourth convolution module, the second output of the third convolution module is connected to the second input of the fourth convolution module, the third output of the third convolution module is connected to the third input of the fourth convolution module, and the fourth output of the third convolution module is connected to the fourth input of the fourth convolution module; The output end of the fourth convolution module outputs the common feature map.

4. The stain detection model training method according to claim 3, characterized in that: The fourth convolution module includes three convolution components connected in sequence, which are respectively recorded as the first convolution component, the second convolution component and the third convolution component; wherein, The first convolution component and the second convolution component each include four convolution branches, which are respectively recorded as a first convolution branch, a second convolution branch, a third convolution branch and a fourth convolution branch. The first convolution branch includes four BasicBlock_SE convolution blocks connected in sequence, the second convolution branch includes four BasicBlock_SE convolution blocks and a 2x upsampling block connected in sequence, the third convolution branch includes four BasicBlock_SE convolution blocks and a 4x upsampling block connected in sequence, and the fourth convolution branch includes four BasicBlock_SE convolution blocks and an 8x upsampling block connected in sequence. The third convolution component includes four convolution branches, a fusion layer and an activation layer, which are respectively recorded as the fifth convolution branch, the sixth convolution branch, the seventh convolution branch and the eighth convolution branch. The fifth convolution branch includes four BasicBlock_SE convolution blocks and an efficient channel attention module connected in sequence. The sixth convolution branch, the seventh convolution branch and the eighth convolution branch all include four BasicBlock_SE convolution blocks, a convolution layer, a batch normalization layer, an upsampling layer and an efficient channel attention module connected in sequence.

5. The stain detection model training method according to claim 4, characterized in that: The BasicBlock_SE convolution block includes a 2D convolution layer, a batch normalization layer, an activation layer, a 2D convolution layer, a compressed excitation attention module, a fusion layer and an activation layer connected in sequence.

6. A stain detection method, characterized in that: The method comprises: Acquire the image to be tested; The image to be tested is input into a pre-trained stain detection model to obtain stain mask information, stain location information and stain classification information in the image to be tested, wherein the pre-trained stain detection model is obtained using the training method of the stain detection model described in any one of claims 1 to 5.

7. A training device for a stain detection model, characterized in that: The stain detection model includes a feature extraction network, a first branch network, and a second branch network, wherein the input end of the feature extraction network is used to input an image, the first output end of the feature extraction network is connected to the input end of the first branch network, and the second output end of the feature extraction network is connected to the input end of the second branch network. The first branch network is used to output predicted stain mask information, and the second branch network is used to output predicted stain location information and stain classification information. The device includes: a data acquisition unit, a model training unit, a function calculation unit, and a parameter adjustment unit. The data acquisition unit, the model training unit, the function calculation unit, and the parameter adjustment unit use the training image set to perform multiple periodic training on the stain detection model. For each training cycle, the following operations are performed: A data acquisition unit, configured to acquire a training image set, wherein the training image set includes an image of the shell appearance; A model training unit, configured to input the shell appearance image in the training image set into a stain detection model to obtain predicted stain mask information, stain location information, and stain classification information; a function calculation unit, configured to obtain a loss function of the stain detection model based on the predicted stain mask information, stain location information, and stain classification information, and label information of the shell appearance image; a parameter adjustment unit, configured to adjust the parameters of the stain detection model according to the loss function of the stain detection model to obtain a trained stain detection model; The feature extraction network includes a convolutional neural network and an extraction candidate frame network. The model training unit is specifically used to input the shell appearance image into the convolutional neural network to obtain a common feature map of the shell appearance image; use the extraction candidate frame network to slide on the common feature map to obtain a common feature map with a stain candidate area; input the common feature map with the stain candidate area into the first branch network to obtain the predicted stain mask information; input the common feature map with the stain candidate area into the second branch network to obtain the predicted stain position information and the predicted stain classification information.

8. The stain detection model training device according to claim 7, characterized in that: The data acquisition unit is specifically used to collect the shell appearance image and pre-process the shell appearance image.

9. The stain detection model training device according to claim 7, characterized in that: The convolutional neural network includes a first convolution module, a second convolution module, a third convolution module and a fourth convolution module connected in sequence. The input end of the first convolution module is used to input an image, the first output end of the first convolution module is connected to the first input end of the second convolution module, and the second output end of the first convolution module is connected to the second input end of the second convolution module: The first output terminal of the second convolution module is connected to the first input terminal of the third convolution module, the second output terminal of the second convolution module is connected to the second input terminal of the third convolution module, and the third output terminal of the second convolution module is connected to the third input terminal of the third convolution module; The first output of the third convolution module is connected to the first input of the fourth convolution module, the second output of the third convolution module is connected to the second input of the fourth convolution module, the third output of the third convolution module is connected to the third input of the fourth convolution module, and the fourth output of the third convolution module is connected to the fourth input of the fourth convolution module; The output end of the fourth convolution module outputs the common feature map.

10. The stain detection model training device according to claim 9, characterized in that: The fourth convolution module includes three convolution components connected in sequence, which are respectively recorded as the first convolution component, the second convolution component and the third convolution component; wherein, The first convolution component and the second convolution component each include four convolution branches, which are respectively recorded as a first convolution branch, a second convolution branch, a third convolution branch and a fourth convolution branch. The first convolution branch includes four BasicBlock_SE convolution blocks connected in sequence, the second convolution branch includes four BasicBlock_SE convolution blocks and a 2x upsampling block connected in sequence, the third convolution branch includes four BasicBlock_SE convolution blocks and a 4x upsampling block connected in sequence, and the fourth convolution branch includes four BasicBlock_SE convolution blocks and an 8x upsampling block connected in sequence. The third convolution component includes four convolution branches, a fusion layer and an activation layer, which are respectively recorded as the fifth convolution branch, the sixth convolution branch, the seventh convolution branch and the eighth convolution branch. The fifth convolution branch includes four BasicBlock_SE convolution blocks and an efficient channel attention module connected in sequence. The sixth convolution branch, the seventh convolution branch and the eighth convolution branch all include four BasicBlock_SE convolution blocks, a convolution layer, a batch normalization layer, an upsampling layer and an efficient channel attention module connected in sequence.

11. The stain detection model training device according to claim 10, characterized in that: The BasicBlock_SE convolution block includes a 2D convolution layer, a batch normalization layer, an activation layer, a 2D convolution layer, a compressed excitation attention module, a fusion layer and an activation layer connected in sequence.

12. A stain detection device, characterized in that: The device comprises: An acquisition module, used for acquiring the image to be tested; A detection module is used to input the image to be tested into a pre-trained stain detection model to obtain stain mask information, stain location information and stain classification information in the image to be tested, wherein the pre-trained stain detection model is obtained using the training method of the stain detection model described in any one of claims 1 to 5.

13. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the training method of the stain detection model according to any one of claims 1 to 5 is implemented, or the stain detection method according to claim 6 is implemented.

14. An electronic device comprising a memory and a processor, wherein a computer program is stored in the memory, wherein: When the computer program is executed by the processor, the training method of the stain detection model according to any one of claims 1 to 5 is implemented, or the stain detection method according to claim 6 is implemented.

Images