A tablet surface defect segmentation method combining feature fusion and attention mechanism

Abstract

The present application belongs to the technical field of image segmentation, and discloses a tablet surface defect segmentation method combining feature fusion and attention mechanism. The specific steps include data preprocessing, constructing an image segmentation network, loading pre-training parameters, and extracting image features by an encoder. The image segmentation network is composed of an encoder and a decoder. The encoder is based on a traditional Unet image segmentation model, and a designed SAM module is selected as a down-sampling block to extract image features. The feature map is input into the decoder through a skip connection, and the features are fused by adopting an attention mechanism to obtain a prediction label map. The test set is predicted based on the tablet defect image segmentation model, and the result of defect area segmentation is output. The model can significantly improve the segmentation performance and effect.

Description

Technical Field

[0001] This invention belongs to the fields of computer vision and tablet segmentation technology, and specifically relates to a method for tablet surface defect segmentation based on a combination of feature fusion and attention mechanisms. Background Technology

[0002] As people's living standards improve, the market demand for pharmaceuticals is increasing, and the pursuit of pharmaceutical quality is becoming increasingly stringent. If defective products enter the market, they severely impact the product's appearance, hindering sales and causing significant economic losses. Furthermore, they affect drug quality, leading to safety hazards during use, such as expired or spoiled medications, potentially causing delays or even worsening of the patient's condition. Therefore, defect detection in pharmaceuticals is crucial for ensuring both the economic benefits of enterprises and the efficacy of drugs. The special nature of pharmaceuticals necessitates supervision and inspection at every stage, from research and development, production, packaging to distribution, making the quality classification and identification of tablets particularly important. How to efficiently classify and remove substandard drugs during production and transportation, and before packaging and warehousing, is an important research direction in the pharmaceutical manufacturing field. Convolutional neural networks (CNNs) possess excellent fault tolerance and self-learning advantages when identifying defective tablet image samples. They are not limited by human subjectivity and error-proneness, nor do they require manual feature selection for detection. Due to their superior performance, CNNs are widely used in defect detection and segmentation. Developing a rapid, accurate, and highly reliable method for detecting defects in pharmaceutical tablets is of research significance and practical value.

[0003] However, due to the large differences in the scale of defects on the surface of the tablets, convolutional networks have difficulty extracting the features of the defects themselves, and small features are easily lost during propagation, resulting in unsatisfactory surface defect segmentation and recognition performance. Summary of the Invention

[0004] The technical problem to be solved by this invention is to provide a method for segmenting tablet surface defects by combining feature fusion and attention mechanisms, which aims to solve the problem of low segmentation accuracy caused by small feature differences extracted by existing encoders.

[0005] To achieve the above objectives, this invention provides a method for segmenting tablet surface defects by combining feature fusion and attention mechanisms, comprising:

[0006] Step 1: Obtain image samples of defects on the surface of the tablet and the corresponding segmentation mask, and preprocess and analyze the data;

[0007] Step 2: Construct an image segmentation network. Based on the traditional Unet image segmentation model, select the designed SAM module as the downsampling block for feature extraction. Input the encoder's feature map into the decoder through skip connections to obtain the SAM-1 image segmentation network.

[0008] Step 3: Based on the SAM-1 image segmentation network, an attention mechanism module is added to the skip connections to obtain the SAM-2 image segmentation network.

[0009] Step 4: Use the data to train the model, calculate the loss between the predicted and true labels of the images, update the model parameters, and obtain the trained segmentation model.

[0010] Step 5: Collect the grayscale image of the surface of the tablet to be detected, and input the grayscale image of the tablet surface into the trained SAM-2 segmentation model to obtain the segmentation label of the tablet to be detected.

[0011] Furthermore, the data preprocessing and analysis in step 1 includes:

[0012] Perform various data augmentation operations on the data, including random rotation, random horizontal and vertical translation, and random scaling.

[0013] Furthermore, step 2, selecting the designed SAM module as the downsampling block for feature extraction, includes: the SAM module consists of a BN layer, an ATT layer, and an FFT layer, as shown below:

[0014] The structure of the ATT layer: The input first undergoes feature transformation through a 1×1 convolutional layer, then is activated by the GELU activation function, and the output serves as the input to the AM module. Finally, a 1×1 convolutional layer transforms the output of the AM module to obtain the final output.

[0015] The structure of an FFT layer: The input x is processed by a 1×1 convolutional layer for feature transformation, followed by feature extraction through a 3×3 depthwise separable convolution. Finally, it is activated by the GELU activation function. Next, a 1×1 convolutional layer is used to transform the output of the 3×3 depthwise separable convolution to obtain the final output.

[0016] The SAM module: The input is standardized by a Batch Normalization (BN) layer, and then features are extracted using two ATT (Active Time Triggers). The outputs of the ATTs are then summed with the inputs of the BN layers (⊕) to obtain an intermediate result. Next, another BN layer is used to standardize the activated result. Finally, an FFT layer is used, and the output of the FFT is summed with the BN inputs of the first layer to obtain the final output.

[0017] Furthermore, in step 3, based on the SAM-1 image segmentation network, the skip connection incorporates an attention mechanism module, which includes: the network consists of four downsampling blocks and four upsampling blocks, and the outputs of the 2nd, 3rd, 4th and 5th blocks in the linear bottleneck layer are used as the outputs of the four stages in the SAM-2 segmentation model network encoder and are incorporated into the skip connection decoder process by the attention SE module.

[0018] Further, step 4 involves using the data to train the model, calculating the loss between the predicted and true labels of the image, updating the model parameters, and obtaining the trained segmentation model, including the following steps:

[0019] Step 4.1 Initialize image segmentation network parameters;

[0020] Step 4.2 sets the training parameters: 300 epochs, 8 batch sizes, 0.95 for the first moment estimate (beta1), 0.9999 for the second moment estimate (beta2), and an initial learning rate (lr) of 1e-4 using the AdamW optimizer.

[0021] Step 4.3 Load training data

[0022] Step 4.4 Iterative Training

[0023] Furthermore, in step 4, the loss between the image predicted label and the real label is calculated, and IOU (Intersection over Union), Dice coefficient, and Jaccard coefficient are used as the evaluation methods for the model.

[0024] This invention enhances feature extraction capabilities by improving the downsampling module and combining it with an attention mechanism for feature fusion. It solves the problems of loss of minute features and loss of transformed features, enabling the network to pay more attention to the defective areas of the tablet, resulting in better segmentation and recognition performance for the surface defects of the tablet. Attached Figure Description

[0025] To more clearly illustrate the technical solutions in this invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only for this invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0026] Figure 1 This is a schematic diagram of the process of the present invention;

[0027] Figure 2 This is a structural diagram of the SAM-2 segmentation model of the present invention;

[0028] Figure 3 This is a structural diagram of the SAM of the present invention;

[0029] Figure 4 This invention includes an ATT and FFT structure diagram;

[0030] Figure 5 This is a structural diagram of the AM module included in this invention. Detailed Implementation

[0031] To better understand the above-mentioned objectives, features, and advantages of the present invention, the present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0032] Many specific details are set forth in the following description in order to provide a full understanding of the invention. However, the invention may also be practiced in other ways different from those described herein, and therefore the scope of protection of the invention is not limited to the specific embodiments disclosed below.

[0033] like Figures 1 to 5 As shown, a method for segmenting defective tablet images combining feature fusion and attention mechanisms includes:

[0034] Step 1: Obtain image samples of defects on the surface of the tablet and the corresponding segmentation mask, and preprocess and analyze the data;

[0035] The data is augmented by performing various data augmentation operations, such as random rotation, random horizontal and vertical translation, and random scaling, to increase the amount of training data.

[0036] Step 2: Construct an image segmentation network. Based on the traditional Unet image segmentation model, select the designed SAM module as the downsampling block for feature extraction. Input the encoder's feature map into the decoder through skip connections to obtain the SAM-1 image segmentation network.

[0037] The SAM module consists of a BN layer, an ATT layer, and an FFT layer, as shown below: The ATT layer structure: The input first undergoes feature transformation through a 1×1 convolutional layer, then is activated by the GELU activation function, and the output serves as the input to the AM module. Finally, a 1×1 convolutional layer transforms the AM output to obtain the final output.

[0038] The structure of an FFT layer: The input x is processed by a 1×1 convolutional layer for feature transformation, followed by feature extraction through a 3×3 depthwise separable convolution. Finally, it is activated by the GELU activation function. Next, a 1×1 convolutional layer is used to transform the output of the 3×3 depthwise separable convolution to obtain the final output.

[0039] The SAM module: The input is standardized by a Batch Normalization (BN) layer, and then features are extracted using two ATT (Active Time Triggers). The outputs of the ATTs are then summed with the inputs of the BN layers (⊕) to obtain an intermediate result. Next, another BN layer is used to standardize the activated result. Finally, an FFT layer is used, and the output of the FFT is summed with the BN inputs of the first layer to obtain the final output.

[0040] like Figure 2 As shown, the relevant formulas are as follows:

[0041] Step 3: Based on the USM-1 image segmentation network, an attention mechanism module is incorporated into the skip connections to obtain the SAM-2 image segmentation network. The step 3, based on the USM-1 image segmentation network, incorporates an attention mechanism module into the skip connections. The network consists of four downsampling blocks and four upsampling blocks. The outputs of the 2nd, 3rd, 4th, and 5th blocks in the linear bottleneck layer serve as the outputs of the four stages in the UNet segmentation model network encoder. An attention SE module is then incorporated into the skip connection decoder process.

[0042] Step 4: Use the data to train the model, calculate the loss between the predicted image label and the real label, update the model parameters, and obtain the trained segmentation model.

[0043] The loss between the predicted and ground truth labels is calculated, and the Intersection over Union (IoU) ratio and the Dice Similarity Coefficient (DSC) are used as evaluation methods for the model. The formula is as follows: Dice = 2 * (|X∩Y|) / (|X| + |Y|)

[0044] IOU = (A∩B) / (A∪B)

[0045] Where A represents the predicted segmentation result; B represents the actual segmentation result;

[0046] Obtaining a trained segmentation model involves the following steps:

[0047] Step 4.1 Initialize image segmentation network parameters;

[0048] Step 4.2 sets the training parameters: 300 epochs, 8 batch sizes, 0.95 for the first moment estimate (beta1), 0.9999 for the second moment estimate (beta2), and an initial learning rate (lr) of 1e-4 using the AdamW optimizer.

[0049] Step 4.3 Load training data

[0050] Step 4.4 Iterative Training

[0051] Step 5: Collect the grayscale image of the surface of the tablet to be detected, and input the grayscale image of the tablet surface into the trained USM segmentation model to obtain the segmentation label of the tablet to be detected.

[0052] The results of comparing it with other commonly used algorithms on the dataset are shown in Table 1.

[0053] Table 1. Comparison of SAM-2 with other algorithms

[0054] mIOU mACC Model parameter count M Unet 83.2 94.1 4.32 Unet++ 83.9 94.9 4.36 ResUnet 83.2 95.8 3.57 SAM-2 87.6 96.2 2.78

[0055] As can be seen from Table 1, the model's IOU and Dice coefficients are slightly higher than those of other algorithms, demonstrating good segmentation capabilities.

[0056] This invention is intended to cover all such substitutions, modifications, and variations that fall within the broad scope of the appended claims. Therefore, any omissions, modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of this invention should be included within the scope of protection of this invention.

Claims

1. A method for segmenting tablet surface defects by combining feature fusion and attention mechanisms, characterized in that, include: Step 1: Obtain image samples of defects on the surface of the tablet and the corresponding segmentation mask, and preprocess and analyze the data; Step 2: Construct an image segmentation network. Based on the traditional Unet image segmentation model, select the designed SAM module as the downsampling block for feature extraction. Input the encoder's feature map into the decoder through skip connections to obtain the SAM-1 image segmentation network. Step 3: Based on the SAM-1 image segmentation network, an attention mechanism module is added to the skip connections to obtain the SAM-2 image segmentation network; Step 4: Use the data to train the model, calculate the loss between the predicted image label and the true label, update the model parameters, and obtain the trained segmentation model. Step 5: Collect the grayscale image of the surface of the tablet to be detected, and input the grayscale image of the tablet surface into the trained USM segmentation model to obtain the segmentation label of the tablet to be detected. Step 2, which selects the designed SAM module as the downsampling block for feature extraction, includes the following: The SAM module consists of a BN layer, an ATT layer, and an FFT layer, as shown below: The structure of the ATT layer: The input first undergoes feature transformation through a 1×1 convolutional layer, then is activated by the GELU activation function, and the output serves as the input to the AM module; finally, the output of the AM module undergoes feature transformation through a 1×1 convolutional layer to obtain the final output. The structure of an FFT layer: The input x passes through a 1×1 convolutional layer for feature transformation, then through a 3×3 depthwise separable convolution for feature extraction; finally, it is activated by the GELU activation function; next, a 1×1 convolutional layer transforms the output of the 3×3 depthwise separable convolution to obtain the final output. SAM module: The input is standardized by a BN layer, and then features are extracted by two ATT layers; Then, the output of ATT is summed with the input of the BN layer by a summation operation ⊕ to obtain an intermediate result; next, the activated result is standardized by a BN layer, and finally, the FFT layer is passed through which the output of the FFT is summed with the BN input of the first layer to obtain the final output; In step 3, based on the SAM-1 image segmentation network, the skip connection incorporates an attention mechanism module, which includes: the network consists of four downsampling blocks and four upsampling blocks, and the outputs of the 2nd, 3rd, 4th and 5th blocks in the linear bottleneck layer are used as the outputs of the four stages in the UNet segmentation model network encoder, which are then combined with the attention SE module in the skip connection decoder process.

2. The method for segmenting tablet surface defects by combining feature fusion and attention mechanisms according to claim 1, characterized in that, Step 1, which involves data preprocessing and analysis, includes: Perform various data augmentation operations on the data, including random rotation, random horizontal and vertical translation, and random scaling.

3. The method for segmenting tablet surface defects by combining feature fusion and attention mechanisms according to claim 1, characterized in that, Step 4 involves using the data to train the model, calculating the loss between the predicted and true labels of the image, updating the model parameters, and obtaining the trained segmentation model. This includes the following steps: Step 4.1 Initialize image segmentation network parameters; Step 4.2 Set the training parameters: the training epoch is 300, the batch size is 8, the exponential decay rate beta1 of the first moment estimate is 0.95, the exponential decay rate beta2 of the second moment estimate is 0.9999, and the initial learning rate lr of the network is set to 1e-4 using the AdamW optimizer. Step 4.3 Load training data; Step 4.4 Iterative training.

Images