An industrial defect small sample data generation and classification method based on staged mask redrawing

Abstract

This invention discloses a method for generating and classifying small sample data of industrial defects based on staged mask redrawing, belonging to the field of industrial visual inspection technology. Addressing the high false negative rate caused by the long-tailed distribution of extremely small industrial defect samples and the problem of mosaic artifacts easily generated by existing generation models, this method first adaptively mines difficult samples based on baseline prediction confidence, and generates a spatial mask for accurately locating defects using a class activation heatmap after truncation, segmentation, and smoothing. Then, a denoising diffusion probability model is used for staged redrawing: in the early stage of sampling, the mask is used for feature-level hard stitching of the hidden states to retain real defect pixels; in the later stage, the mask is removed, and a global self-attention mechanism is used for autonomous smoothing updates. Finally, the generated high-fidelity samples are appended to the extremely small sample set and input into the downstream network for end-to-end training using dynamic class weighting. This invention completely eliminates abrupt changes in illumination and harsh artifacts in the generated images, significantly broadens the decision boundary of the classifier, and reduces the actual false negative rate.

Description

Technical Field

[0001] The invention relates to the fields of industrial visual inspection, deep learning, and computer vision intelligent manufacturing, specifically to a method for augmenting and intelligently classifying industrial defect image data under long-tailed distribution and extremely small sample conditions. Background Technology

[0002] Deep learning technology has been widely used in industrial visual inspection, but it has a fatal weakness: it is extremely dependent on massive amounts of high-quality training data. However, in real factory production lines, severe surface defects in steel are extremely rare, resulting in very few images of certain categories being collected, causing a long-tail effect of extremely unbalanced data distribution. Detection models trained under this "starved" state are prone to serious false negatives in real-world applications. In industry, false negatives mean that defective products will directly enter the next process or the market, which is absolutely unacceptable.

[0003] To address the problem of data scarcity, existing technologies typically employ data augmentation methods:

[0004] (1) Traditional geometric enhancement methods

[0005] Such methods merely transform the existing image pixels and cannot create new defect morphologies or industrial backgrounds out of thin air, thus failing to help classification networks broaden the decision boundaries of their learning.

[0006] (2) Mainstream generative models

[0007] Such methods are prone to pattern collapse under the constraint of extremely small sample sizes. When attempting to generate new defects, they often disrupt the original normal industrial background texture, resulting in images with obvious "awkward patchwork" or "sudden changes in lighting." These unrealistic and physically incompatible generation results actually introduce harmful noise into the downstream classification network, which in turn interferes with the normal convergence of the model. Summary of the Invention

[0008] To address the problems of tiling artifacts, background texture corruption, and high false negative rates in existing generative models with extremely small sample sizes, this invention proposes a method for augmenting and classifying industrial defects using extremely small sample data based on staged mask redrawing. This method delves into the underlying ODE solver of the generative model, introducing "same-frequency aligned noise" and a staged "early release" strategy, achieving the dual objectives of strictly preserving the defect subject and allowing background features to evolve naturally.

[0009] To achieve the above objectives, this invention proposes a method for augmenting and classifying minimal sample data of industrial defects based on staged mask redrawing, which includes the following steps:

[0010] S1: Mining of difficult samples based on confidence ranking.

[0011] Input defect images and pre-trained baseline classification networks to obtain baseline prediction confidence scores for samples; rank samples in ascending order based on the confidence scores, adaptively mine the difficult samples that the model is most likely to misclassify, and construct a difficult anchor subset.

[0012] Specifically, given a pre-trained baseline classification network For those belonging to the category The original input image The baseline classification network is used to calculate its confidence score for the true category. The calculation formula is:

[0013]

[0014] in, For the corresponding defect category label, Indicates the baseline classification network parameters Next, input image Predicted as category The probability. For each category, all samples are sorted by confidence level. Sort in ascending order and truncate the results to the lowest confidence level. These samples constitute a subset of the difficult anchor points for this category.

[0015] S2: Extract deep features for difficult samples and generate spatial masks.

[0016] The deep feature gradient information of the difficult samples is extracted using a baseline classification network to generate a class activation heatmap; the heatmap is then smoothed and segmented to obtain a spatial mask for accurately locating the core defect region.

[0017] Specifically, the spatial mask is obtained through the feature extraction and image processing modules. The calculation formula is:

[0018]

[0019] in, Grad-CAM heatmap representing the difficult sample; Indicates The threshold-based truncation / segmentation operation; The kernel size is represented by a parameter. The determined Gaussian smoothing function. The spatial mask. The range of values ​​is For background samples, directly set .

[0020] S3: Inject initial noise that is globally aligned with the same frequency.

[0021] Globally fixed alignment noise is injected into the real input image, and the theoretical state at any time step in the forward process is calculated to ensure that the masked region and the unmasked region share the same initial noise frequency distribution, which serves as the physical reference for the subsequent generation process.

[0022] Specifically, in time step During the forward pass, the input image at any time Theoretical state The calculation formula is:

[0023]

[0024] in, Represents the actual input image; This represents the injected global fixed alignment noise, and it follows a standard normal distribution. .

[0025] S4: Staged mask redrawing based on Heun ODE solver.

[0026] Set the total sampling time steps and time threshold The updated hidden state The data is fed into the Heun ODE solver for iterative updates. Within each sampling time step, the hidden state is updated according to the following phased formula:

[0027]

[0028] in, This represents the element-wise multiplication operation; This represents the hidden state generated at the current time step;

[0029] (1) If the current sampling time step is less than the time threshold ( Then, the mandatory constraint phase begins: utilizing the aforementioned spatial mask. Perform feature-level hard stitching update of the hidden state. Force the retention of real defective pixels;

[0030] (2) If the current sampling time step is greater than or equal to the time threshold ( Then, it enters the natural smoothing stage: cancel the spatial mask constraint and directly set... The hidden state is updated globally autonomously and smoothly using a global self-attention mechanism.

[0031] After completing the current time step, update the hidden state and time step until all sampling time steps are completed.

[0032] S5: Map the output high-fidelity image and build an expanded dataset.

[0033] After completing all sampling time steps, the mask overlay is discarded and linear inverse normalization is performed to map and output a high-fidelity generated image, thus completing the local background redrawing of a single difficult sample; the filtered generated image set is then... Added to the original minimal training set Build an expanded dataset The merging formula is:

[0034]

[0035] in, This represents the augmented dataset ultimately used for training. This represents the original, minimal training set. This represents the generated expanded sample set.

[0036] S6: Use class-weighted loss to train the classification network using the original samples.

[0037] The expanded dataset The input is fed into a classification network, using a class-weighted cross-entropy loss function. End-to-end iterative training is performed to obtain the final intelligent classification model for industrial defects. The formula for calculating the loss function is as follows:

[0038]

[0039] in, This represents the total number of samples in the expanded dataset; Represents the input image. Represents the actual defect category label; Indicates the classification network in terms of parameters Below is the input image Predicted as true label The probability of.

[0040] The For categories The adaptive penalty weight is calculated using the following formula:

[0041]

[0042] in, This represents the total number of samples in the expanded dataset. Represents the current category The total number of samples in the expanded dataset.

[0043] The beneficial effects of this invention are mainly reflected in:

[0044] (1) An adaptive hard sample mining and targeted generation augmentation strategy based on confidence ranking is proposed. By obtaining the prediction confidence of the baseline classification network and sorting it in ascending order, the hard anchor points that the model is prone to misclassification are accurately screened. Compared with the traditional indiscriminate full data augmentation method, this method avoids the "gradient dilution" problem caused by generating a large number of redundant simple samples, and accurately allocates the limited generation computing resources to the vicinity of the decision boundary, thus efficiently and specifically expanding the feature space of complex categories.

[0045] (2) A staged mask redrawing mechanism based on the Heun ODE solver was designed, and global frequency-aligned initial noise was innovatively introduced into the forward integration path. By implementing a forced hard-joining constraint of the spatial mask in the early stage of sampling and implementing a natural smooth "early release" strategy of canceling the mask in the later stage, the physical form of the original real defect subject can be strictly preserved, while promoting the autonomous evolution of the background global context. Compared with traditional local masking or image tiling generation techniques, this invention completely eliminates high-frequency feature conflicts, illumination abrupt changes and edge artifacts from the physical generation level, and greatly improves the high fidelity and visual continuity of the augmented image.

[0046] (3) Combining the characteristics of extremely small samples and extreme long-tail distribution in industrial quality inspection scenarios, an incremental joint training method using class-weighted cross-entropy loss is proposed. The constructed expanded dataset and the original samples are jointly input into the network. By introducing an adaptive penalty weight inversely proportional to the number of class samples into the loss function, the classification model is forced to increase its attention to the extremely scarce minority classes. Compared with conventional model optimization techniques, this invention significantly fills the feature gap of extremely small samples, broadens the decision boundary of the model, greatly reduces the serious false negative rate caused by the long-tail effect in practical industrial applications, and comprehensively improves the overall robustness and recall of the intelligent classification model. Attached Figure Description

[0047] Figure 1 This is a flowchart of a defect image data augmentation and classification method based on staged mask redrawing provided by an embodiment of the present invention;

[0048] Figure 2 This is a schematic diagram illustrating the effect of difficult sample feature localization and staged mask redrawing to generate high-fidelity data variants according to an embodiment of the present invention;

[0049] Figure 3 This is a comparison diagram of the qualitative visual effects of different data augmentation and generation methods provided in an embodiment of the present invention on various defects;

[0050] Figure 4This is a semantic segmentation task extension task generation graph provided by an embodiment of the present invention based on a private dataset.

[0051] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments.

[0052] Example 1: Application and Performance Verification Based on Downstream Image-Level Classification Task

[0053] This embodiment provides a specific implementation of a method for augmenting and classifying small sample data of industrial defects based on staged mask redrawing, focusing on verifying the generalization performance and application effect of the invention in the downstream image-level classification task. The specific implementation steps of this embodiment are as follows:

[0054] (1) Experimental dataset and network benchmark settings

[0055] This embodiment uses a real industrial steel surface defect dataset for testing and verification. This dataset exhibits an extreme long-tail distribution and contains multiple defect categories. Initially, each defect category has a very small sample size (e.g., a single category contains only 181 independent original defect samples).

[0056] In this embodiment, the downstream main task classification network uses the standard ResNet-18 network. The AdamW optimizer is used during network training, with an initial learning rate of [missing information]. The weight decay parameter is set to .

[0057] (2) Phased mask redrawing and data augmentation

[0058] 1. Transition between coarse and fine mask generation:

[0059] For the original minimal sample set Each defect image in (in (For the corresponding category label), first obtain its initial defect area or sensitive area hints. The system utilizes geometric transformation operators and smoothing operators. The process is performed to generate a series of soft masks. The specific calculation formula is as follows:

[0060]

[0061] In the formula, Indicates The threshold-based truncation segmentation operation is used to spatially widen, rotate, or morphologically perturb the defect region indication; The standard deviation is expressed as The Gaussian blurring operator is used to eliminate step abrupt changes at the mask edges, thereby ensuring that the generated soft mask satisfies... This ensures a smooth transition between old and new pixels at the boundary.

[0062] 2. Diffusion redraw based on pure Transformer architecture:

[0063] This invention constructs a denoising diffusion probability model based on the pure Transformer architecture (Just Image Transformers, JiT) to replace the traditional U-Net architecture that highly relies on local convolutions.

[0064] During the redrawing stage, the input defective image is divided into several non-overlapping image patch sequences. The long-distance spatial dependencies between all image patches are directly modeled through the global self-attention mechanism, thereby breaking the local translation invariance limitation of traditional convolutional networks.

[0065] Combined with the mask generated in step 1 This method utilizes a diffusion redrawing mechanism in latent space or pixel space to conditionally redraw non-masked regions (background or specific areas). The resulting augmented defect image... The following latent variable mixing formula is obtained:

[0066]

[0067] In the formula, This represents the pixel-wise multiplication operator; This indicates a denoising diffusion generator based on the JiT architecture; This represents a random noise sample of the input; This represents the augmentation factor for a single sample. This redrawing mechanism, while preserving the core morphological features of the original defects, can efficiently reconstruct high-frequency textures (such as discrete oxide scale and narrow scratches). Furthermore, by eliminating the interference of invalid labels as functional background classes during the generation process, the redrawing results more closely resemble the physical distribution of a real industrial production line. By iteratively redrawing each extremely small sample defect category, a high-quality augmented sample set is ultimately constructed. .

[0068] 3. Dataset database merging:

[0069] The generated expanded sample set The augmented dataset is created by directly writing images in a specified image format (such as PNG) to the storage hard drive, and then combining it with the original minimal sample set. Perform a database merge operation to generate a complete expanded dataset. .

[0070] (3) Classification model training and class-weighted loss optimization

[0071] The expanded dataset The input is fed into the ResNet-18 classification network for iterative training.

[0072] To eliminate the severe interference caused by long-tailed distributions and extremely small samples on the decision boundary of the classification network, the system calculates the total number of samples of each class in the current training batch or the entire dataset in real time when calculating the loss. The system dynamically generates penalty weights for each defect category based on the sample distribution. The calculation formula is as follows:

[0073]

[0074] In the formula, This represents the total number of samples within the current statistical scope. The penalty weight is then applied. The gradient is calculated by substituting the values ​​into the category-weighted cross-entropy loss function, and backpropagation is performed to update the network weights. This mechanism automatically increases the penalty for defect categories with extremely small sample sizes during training, guiding the model to pay balanced attention to various types of defects. Finally, the converged model weights (in .pth format) are exported for intelligent defect detection in actual production lines.

[0075] 4. Analysis of the experimental results and verification of the technical advantages of the classification task

[0076] To objectively evaluate the effectiveness of the method of this invention, a rigorous quantitative and qualitative comparison was conducted on an industrial steel surface defect classification test set, comparing it with existing mainstream data augmentation baseline methods (including traditional geometric enhancement Mixup, ACGAN based on generative adversarial networks, and the standard conditional diffusion probability model DDPM). All experiments followed the principles of "equal budget, equal network, equal evaluation," and uniformly adopted an extremely small sample benchmark (only 181 images per class).

[0077] The evaluation metrics for the classification task are macro-F1 score, overall accuracy, precision (P), and recall (R).

[0078] Table 1 compares the performance indicators of different models of this invention;

[0079] Table 1

[0080]

[0081] Experimental results show that:

[0082] (1) In terms of model classification and prediction performance:

[0083] The method of this invention significantly improves the classification performance of extremely small samples under long-tailed distribution.

[0084] 1. Mixup can only perform simple pixel-level linear overlay, which cannot enrich semantic information, and its Macro-F1 is only 0.7192.

[0085] 2. Although ACGAN increases the recall rate in category 1 (C1) through nonlinear feature perturbation, it causes serious false positives, resulting in a sharp drop in precision to around 0.40. This "artificially high" recall rate has no practical value in demanding industrial inspection scenarios.

[0086] 3. In the absence of hard constraints from the underlying spatial mask, standard DDPM's global blind denoising can easily cause tiny defect features to be "washed away" as noise, resulting in a Macro-F1 of only 0.7556.

[0087] 4. In comparison, the classification network trained using the method of this invention achieves optimal results in both Macro-F1 (0.7715) and overall accuracy (0.88). Particularly in the difficult sample class (C1), this invention not only boosts the recall to 0.93 but also successfully curbs the spread of false positives, maintaining a stable precision of 0.49, ultimately achieving an F1 score of 0.64 (far exceeding ACGAN's 0.55). Similarly, in the C0 class, this invention maintains extremely high precision (0.69) while improving recall. The data fully demonstrates that the data generated by this invention is not meaningless noise interference but accurately supplements the highly discriminative features of difficult samples, broadening the decision boundaries of the classifier.

[0088] (2) Regarding data generation fidelity and computational efficiency:

[0089] The method of this invention exhibits significant advantages in image quality and augmentation efficiency.

[0090] 1. Traditional ACGAN-generated images suffer from severe mode collapse and high-frequency grid artifacts, with an average FID value as high as 280.

[0091] 2. Traditional ACGAN-generated images suffer from severe mode collapse and high-frequency grid artifacts, with an average FID (Fréchet Inception Distance) as high as 280. This indicates that the classifier is highly likely to get caught in shortcut learning, misclassifying artifacts as classification cues.

[0092] 3. Standard DDPM not only faces high time costs (training takes about 40 hours for the same number of rounds), but also generates samples with obvious "overexposure" and low contrast, failing to effectively fit the complex local lighting changes in real industrial scenes, with an FID of 129.

[0093] 4. The method of this invention benefits from the global receptive field and staged ODE resampling mechanism of the JiT architecture, with a single training session taking only about 13 hours, and computational efficiency nearly 3 times higher than the standard DDPM. Simultaneously, this invention optimizes the FID mean to 90, eliminating mesh artifacts and abrupt lighting changes, achieving seamless physical fusion between the original defective subject and the generated background pixels.

[0094] Example 2: Extended Application and Effect Verification Based on Downstream Semantic Segmentation Tasks. To further verify the cross-task generalizability of the proposed minimal sample data synthesis strategy and its practical value on real industrial production line data, this example extends the test benchmark from image-level classification tasks to more refined pixel-level semantic segmentation tasks.

[0095] 1. Experimental Dataset and Network Setup

[0096] This embodiment uses a real industrial steel surface defect dataset for extended validation. This dataset includes background images (category 0) and five real defect categories (categories 1 to 5). The dataset exhibits an extreme long-tail distribution, with the training set for the extremely rare category (category 5) containing only 18 images, used to realistically simulate this extremely rare defect type in industrial settings.

[0097] In terms of model construction, this embodiment uses a standard convolutional neural network (preferably ResNet-34) as the basic segmentation network, and employs weighted cross-entropy and Dice Loss for joint optimization. To verify the effectiveness of this invention, the benchmark model and the method of this invention maintain the same model structure and hyperparameters, and only the augmented synthetic images generated by the method of this invention are introduced into the training set for comparison.

[0098] 3. Quantitative Results and Effect Analysis

[0099] The generation effects of the baseline model and the model extended using the method of this invention on defect segmentation of various categories are as follows: Figure 4 As shown in Table 2, the performance indicators are compared.

[0100] Table 2 compares the performance metrics of semantic segmentation task extension tasks;

[0101] Table 2

[0102]

[0103] Analysis of the experimental data in Table 2 reveals the following significant beneficial effects of this invention:

[0104] Overall segmentation performance and anti-interference capability are significantly improved: After incorporating the images synthesized by this invention, the average IoU across all categories increased from 0.6436 to 0.6662. More importantly, the average accuracy significantly improved from 0.7468 to 0.7959. This indicates that the high-fidelity background features generated by this invention effectively suppress the "feature illusion" of the model in complex industrial backgrounds, greatly reducing the probability of misreporting normal backgrounds as defects.

[0105] Successfully overcame the challenge of detecting extremely small samples: For the most difficult category (5, non-metallic inclusions) with only 18 samples, where strong baselines were already difficult to overcome, the method of this invention achieved a leap in mean IoU from 0.5299 to 0.5990 (nearly 7%). This directly proves that the targeted redrawing strategy of this invention can accurately fill in the morphological and feature gaps of rare defects.

[0106] Adaptive rebalancing of the model's decision boundary: Although the IoU metric for common, simple classes (such as class 1) showed a slight decline, this reflects that the model has successfully overcome its "overfitting" bias towards the majority class. The model redistributed previously redundant attention weights to more challenging and industrially valuable hard classes such as C3 and C5. This strategy of sacrificing a small portion of performance in regular classes for qualitative improvement in hard classes makes the intelligent detection model more robust to real-world, complex industrial scenarios.

[0107] In summary, the data augmentation method proposed in this invention can not only effectively solve the problem of missed detection in classification tasks, but also demonstrates excellent feature enhancement and cross-task generalization capabilities when dealing with more complex pixel-level localization (semantic segmentation) tasks, and has extremely high practical industrial application value.

[0108] Finally, it should be noted that the above descriptions are merely some embodiments of the present invention. Obviously, the present invention is not limited to the above embodiments, and many variations are possible. All variations that can be directly derived or conceived by those skilled in the art from the content disclosed in this invention should be considered within the scope of protection of this invention.

Claims

1. A method for generating and classifying small sample data of industrial defects based on staged mask redrawing, characterized in that: The method includes the following steps: S1, obtain the baseline prediction confidence score of the input image, and adaptively mine difficult samples that the classification network is prone to misclassification based on the confidence score to construct a difficult anchor subset; S2, using a pre-trained baseline classification network to extract deep feature gradient information of the difficult samples to generate class activation heatmaps, and performing spatial segmentation and smoothing to obtain spatial masks for accurately locating core defect regions; S3 injects globally fixed, same-frequency aligned noise into the input image and calculates the noise mixing state at any time step during the forward denoising process. S4, use the denoising diffusion probability model to redraw the mask, and set the total sampling time step and time threshold: in the early stage of sampling, use the spatial mask to perform feature-level hard stitching update of the hidden state to retain real defect pixels; in the later stage of sampling, cancel the spatial mask constraint and use the global self-attention mechanism to perform autonomous smooth update of the hidden state. S5, after completing all sampling time steps, map the high-fidelity generated image and append the generated expanded sample set to the original minimal sample training set to construct the expanded dataset; S6. Input the expanded dataset into the downstream classification network and perform end-to-end training using the category-weighted cross-entropy loss function to obtain the final intelligent classification model for industrial defects.

2. The method of claim 1, wherein, Before step S1, the method also includes obtaining the original minimal sample training set: The original industrial defect image is segmented into independent image slices using a non-overlapping sliding window cropping strategy. Simultaneously decode the run-length encoding into a slice-level binary mask; Hard filtering of the image slices is performed based on a set mask area threshold to construct a single-label classification dataset for training.

3. The method of claim 1, wherein, In step S1, the confidence score is calculated The formula is: For each category All samples were sorted by confidence level Sort in ascending order and truncate the results to the lowest confidence level. These samples constitute the subset of the difficult anchors for this category.

4. The method of claim 1, wherein, In step S2, the spatial mask The formula for calculating the spatial mask is: wherein a class activation heat map representing the difficult sample; denotes a thresholded segmentation operation with a threshold value; denotes a Gaussian smoothing function with a standard deviation of .

5. The method of claim 1, wherein, In step S4, the hidden states are fed into the ordinary differential equation solver The update is performed according to the following piecewise formula: In the formula, The space mask; This is the current sampling time step; The set time threshold; This represents the element-wise multiplication operation; This refers to the noise-mixing state; The hidden state predicted at the current time step.

6. The method of claim 1, wherein, In step S6, the category-weighted cross-entropy loss function The calculation formula is: Wherein, the adaptive penalty weight for each defect category The calculation formula of the adaptive penalty weight is: wherein and both represent the total sample size of the augmented dataset; representing the current class total number of true samples in the augmented dataset.

Images