Dynamic Label Injection for Imbalanced Industrial Defect Segmentation

TL;DR

This work tackles imbalanced multi-class semantic segmentation in industrial defect datasets by introducing Dynamic Label Injection (DLI), an online augmentation pipeline that balances each training batch by injecting augmented defects into defect-free images. The injection combines Poisson-based seamless cloning and cut-paste to ensure realistic blending and diverse appearances, with a formal algorithm that selects the rarest class in the current batch and injects its defects until uniformity is achieved. Experiments on the Magnetic Tiles dataset show that DLI substantially improves mean IoU over standard balancing losses across lightweight UNet backbones, and ablations confirm the complementary value of combining both injection techniques. The method demonstrates robustness in weakly supervised settings and offers a practical, open-source solution for industrial segmentation under severe class imbalance.

Abstract

In this work, we propose a simple yet effective method to tackle the problem of imbalanced multi-class semantic segmentation in deep learning systems. One of the key properties for a good training set is the balancing among the classes. When the input distribution is heavily imbalanced in the number of instances, the learning process could be hindered or difficult to carry on. To this end, we propose a Dynamic Label Injection (DLI) algorithm to impose a uniform distribution in the input batch. Our algorithm computes the current batch defect distribution and re-balances it by transferring defects using a combination of Poisson-based seamless image cloning and cut-paste techniques. A thorough experimental section on the Magnetic Tiles dataset shows better results of DLI compared to other balancing loss approaches also in the challenging weakly-supervised setup. The code is available at https://github.com/covisionlab/dynamic-label-injection.git

Figures (8)

• Figure 1: Overview of Dynamic Label Injection algorithm. $\mathcal{D}_{train}$ contains defect-free and defective samples. During training, defects are dynamically sampled, augmented, and injected into defect-free images to balance the batch $\mathcal{B}^{*}$, which is then fed into a multi-class defect segmentation neural network.
• Figure 2: Magnetic Tiles dataset class imbalanced distribution.
• Figure 3: IoU scores for each defect averaged over 5 different seeds. An horizontal red line highlights the gap between DLI and competitors.
• Figure 4: Qualitative samples for each class of defects in the Magnetic Tiles dataset.
• Figure 5: The plot shows how IoU varies with the percentage of data in the training set, using DLI and WCE loss. The left chart displays absolute IoU values, while the right chart shows the performance drop compared to training with the full dataset.