A method for infrared small target segmentation based on weak supervision

Abstract

The present invention discloses a method for segmenting infrared small targets based on weak supervision. The input of the present invention is only the infrared small target image and the infrared small target bounding box data, and the final output result is the segmentation mask of the small target in the infrared small target image. The present invention includes model pre-training, model loading, teacher model generation of pseudo mask, student model training and dual-model EMA, wherein the teacher model generation of pseudo mask, student model training and dual-model EMA need to be repeated for multiple iterative training. Based on the use of weak supervision loss PS‑Loss, the present invention introduces an adaptive learning rate adjustment mechanism to dynamically adjust the learning rate according to the performance of the model on the validation set. In the weak supervision training process of infrared small targets, an EMA-based strategy, i.e., exponential moving average, is adopted. By maintaining two models, one for training and the other for obtaining slight changes in the infrared small target image, the instability and fluctuation of the model during training are successfully reduced.

Description

Technical Field

[0001] The present invention provides an infrared small target segmentation method based on weak supervision. Background Art

[0002] The scarcity of natural image datasets has always been one of the reasons that hinder the further improvement of the performance of natural image segmentation algorithms. Similarly, this problem also plagues infrared small target segmentation algorithms. This challenge has prompted researchers to seek more adaptive methods to deal with the problem of data scarcity. In this context, weakly supervised algorithms have become one of the solutions. By utilizing data with only coarse labels or less annotated information, weakly supervised algorithms can train models in the absence of sufficient annotations. In the field of infrared small target segmentation, the introduction of weakly supervised algorithms has brought new possibilities for overcoming data scarcity. By redesigning the loss function and training strategy, these algorithms can effectively utilize a small amount of annotated information to achieve accurate segmentation of infrared small targets. Therefore, weakly supervised algorithms have become one of the hottest research directions in infrared small target segmentation.

[0003] Currently, no work in the field of infrared small target segmentation has used weak supervision to address the dataset scarcity issue. To fill this research gap, this paper addresses this issue for the first time from a weakly supervised perspective, introducing a weakly supervised image segmentation method called BoxInf (Box-based weakly supervised segmentation method for infrared small targets) designed for the characteristics of infrared small targets. Summary of the Invention

[0004] The purpose of the present invention is to address the shortcomings of the existing technology and provide a weakly supervised infrared small target segmentation method. The present invention only inputs the infrared small target image and the infrared small target bounding box data, and the final output is the segmentation mask of the small target in the infrared small target image. The present invention includes model pre-training, model loading, pseudo-mask generation by the teacher model, student model training, and dual-model EMA. The teacher model pseudo-mask generation, student model training, and dual-model EMA training require multiple iterations of training.

[0005] The technical solutions adopted by the present invention to solve the technical problems are as follows:

[0006] Step 1: Use the infrared small target image and the corresponding bounding box mask to pre-train the model in the PS-Loss loss function.

[0007] Step 1: Copy the parameters of the pre-trained model to the teacher model and the student model. Fix the parameters of the teacher model for use in subsequent steps.

[0008] Step 2: Generate a pseudo mask for the infrared small target image through the teacher model.

[0009] Step 3: Compare the pseudo mask with the bounding box mask and remove the pseudo mask outside the bounding box. This is because the small infrared target has been pre-set to be within the bounding box.

[0010] Step 4: Use the bounding box mask from step 1 and the pseudo mask processed in step 3 as guidance labels for the loss function to control the student model training and obtain the output of the student model. When training the student model through the loss function, the bounding box mask uses the PS-Loss loss function, while the pseudo mask uses the general image segmentation loss function DiceIoU-Loss loss function.

[0011] Step 5: Input the student model and the teacher model into the EMA algorithm to update the teacher model, then return to step 3 and perform multiple repeated iterative training. The final student model is used as the infrared small target segmentation model, and the prediction mask output by the final student model is used as the final output for evaluation.

[0012] The beneficial effects of the present invention are as follows:

[0013] Based on the use of a weakly supervised loss (PS-Loss), this paper introduces an adaptive learning rate adjustment mechanism that dynamically adjusts the learning rate based on the model's performance on the validation set. This helps the model better adapt to different data distributions and scenarios, improving generalization. During weakly supervised training of infrared small targets, a strategy based on the exponential moving average (EMA), or exponential moving average, is employed. By maintaining two models—one for training and the other for capturing subtle changes in infrared small target images—this strategy successfully reduces model instability and fluctuation during training. This strategy not only improves training robustness but also accelerates convergence. Furthermore, data augmentation techniques, including random rotation, scaling, and horizontal flipping, are introduced during training to increase the model's adaptability to different scales and angles. Color enhancement is also applied to the student model. These methods help improve the model's generalization performance, making it more robust in practical applications. The proposed framework for weakly supervised training of infrared small targets comprehensively considers multiple aspects, including loss function design, learning rate adjustment, EMA strategy, and data augmentation, to achieve better performance and robustness. Experiments have verified the effectiveness of this framework on various datasets, achieving satisfactory performance improvements.

[0014] In summary, this paper utilizes infrared small target bounding box masks as supervision, combines out-of-box loss, in-box loss, and a dual-model-based EMA iterative algorithm, and provides a more cost-effective training method for infrared small target segmentation. Without the need for large-scale annotation, the model can acquire knowledge from limited weak labels, resulting in excellent performance in infrared small target segmentation tasks. Experimental results show that it successfully overcomes the tedious and expensive problems of manual annotation methods, providing a more efficient and sustainable training method for infrared small target segmentation tasks. BRIEF DESCRIPTION OF THE DRAWINGS

[0015] Figure 1 This is the overall framework of the BoxInf weak supervision algorithm for infrared small targets.

[0016] Figure 2 The scope of loss outside the frame and loss inside the frame.

[0017] Figure 3 is the mapping between the real mask and the bounding box mask.

[0018] Figure 4 This is the flowchart of the out-of-box loss function Project-Loss.

[0019] Figure 5 The principle of edge shape generation for small infrared targets.

[0020] Figure 6 Loss edge points and internal points within the box.

[0021] Figure 7 This is the Shape-Loss flow chart for the in-box loss.

[0022] Figure 8 False positive rate of segmented data with different thresholds on the infrared small target dataset.

[0023] Figure 9 Convert the infrared small target dataset into the dataset required for weak supervision.

[0024] Figure 10 Compare the weakly supervised algorithm with other algorithms. DETAILED DESCRIPTION

[0025] The present invention will be further described below with reference to the embodiments.

[0026] The weak supervision loss function PS-Loss of infrared small targets consists of the projection loss Project-Loss outside the frame and the shape loss Shape-Loss inside the frame. Figure 2As shown in the figure, according to the different scopes of the two losses, the goal of the out-of-box loss is to confine the small infrared target within the bounding box, while the purpose of the in-box loss is to adjust the shape of the small target within the bounding box. Specifically, the out-of-box projection loss maps the maximum value of the predicted image to the x-axis and y-axis, and calculates the loss values ​​on the x-axis and y-axis by comparing the bounding box mask. In contrast, the in-box loss uses the prior knowledge of the infrared small target image as the label value, obtains the edge pixels and internal pixels of the infrared small target image as labels through the directional difference convolution kernel, and then calculates the shape adjustment loss. The final infrared small target weak supervision loss L ps Defined as:

[0027] L ps =L project +λL shape (1)

[0028] Among them L project is the out-of-frame projection loss, with a weight of 1; L shape The loss is the in-box shape loss, with a default weight of λ, which was set to 0.1 in the experiment. The weakly supervised framework adapts the loss function without modifying the existing small infrared target detection model. This loss function design enables the training framework to effectively guide the model in learning accurate localization and fine segmentation of small infrared targets, thereby improving model performance.

[0029] The specific operation of Project-Loss outside the frame is as follows Figure 3 As shown in , the predicted mask output by the infrared small target segmentation model needs to be mapped to the x-axis and y-axis. At the same time, the bounding box mask of the infrared small target image also needs to be mapped in the same way. Subsequently, DiceIoU-Loss is introduced for calculation. This loss function is used to measure the accuracy and position alignment of the small target predicted by the model within the box, as shown in Figure 4 As shown in the figure, DiceIoU-Loss effectively guides the model to ensure that the predicted mask output by the model is as consistent as possible with the actual position and shape by comparing its similarity with the actual label on the x- and y-axes. The introduction of this loss function further enhances the method's performance in small object localization and segmentation accuracy. By comprehensively considering both position and shape information, this method demonstrates superior performance in infrared small object segmentation tasks.

[0030] The main purpose of the out-of-box projection loss Project-Loss is to map small targets into the bounding box mask. By ensuring that the bounding box mask map matches the predicted mask map on the x-axis and y-axis, the model has the ability to distinguish small targets from the background. Where m represents the bounding box mask and b represents the predicted mask output by the student model. The projection process obtains the corresponding projection m on the x-axis and y-axis by calculating the maximum value of the bounding box mask on the y-axis and x-axis.x and m y , as shown in formula (2). Similarly, the corresponding projections b of the predicted mask on the x-axis and y-axis can be obtained x and b y , as shown in formula (3). Then, the DiceIoU-Loss loss between the bounding box mask projection and the predicted mask projection is calculated. Formula (4) is the loss function formula of the out-of-box loss Project-Loss. The formula is as follows:

[0031] Project x (m)=max y (m) = m x , (2)

[0032] Project y (b) = max x (b) = b x , (3)

[0033] L Project =L DiceIoU (m x ,b x )+L DiceIoU (m y ,b y ), (4)

[0034] The loss function is to minimize the difference between the bounding box mask and the predicted mask. y (m) represents the mapping of the bounding box mask on the y-axis, and max x (b) shows the mapping of the predicted mask on the x-axis.

[0035] Shape-Loss uses a directional difference convolution kernel to extract the shape features of small infrared targets and pre-sets labels based on the shape features to generate prior labels. The core of this idea is to use the directional difference convolution kernel method to capture the shape features of small infrared targets and distinguish between internal and external pixels within the bounding box. Figure 5 This process is demonstrated in this paper. By using this convolution kernel, shape feature labels can be pre-set, guiding the label generation process. This approach helps more accurately guide the model to learn the shape information of small infrared targets, improving the model's performance in terms of in-frame loss. This design enables the model to better adapt to the shape changes of small infrared targets, further enhancing the accuracy and robustness of segmentation.

[0036] Specifically, for the infrared small target image I∈R w*h , use the x-axis and y-axis direction differential convolution kernels v' and h' to extract the shape feature map U of the infrared small target in two directions x and U y, and then use the square sum of the two variables to generate the shape U, through the threshold setting, the larger and smaller values ​​are set as edge pixels and internal pixels respectively, and the edge shape graph U∈R is generated by the edge pixels and internal pixels w*h , guiding the model to perform training optimization. The formula for obtaining the edge shape graph is as follows:

[0037]

[0038] For the predicted value, it needs to be mapped to an edge value or an internal value. The prior knowledge of this idea is that it is generally believed that a small target is a whole with no ring inside and is stable. The predicted values ​​of two pixels in the x-axis or y-axis direction of the pixel point are the same to classify them. If they are the same, the pixel point is an internal pixel, and if they are different, the pixel point is an external pixel. Figure 6 As shown in the figure, the gray area is the background, the white area is the target, and after edge differentiation, the red area is the internal pixel and the blue area is the external pixel. By obtaining the values ​​of the left and right ends of a pixel and calculating their difference, it is determined whether the pixel is an internal pixel or a boundary pixel of the infrared small target. After obtaining the prediction image, for each pixel, it is necessary to calculate the similarity between it and the two predicted values ​​on the left and right to determine whether it is an internal pixel on the x-axis. By calculating the consistency of its neighbor labels, if it is higher than the set threshold λ, it is judged as an internal pixel, and if it is lower than the threshold α, it is judged as an external pixel. Pixels with moderate thresholds are relatively blurry and should be discarded in the loss calculation.

[0039] The in-frame shape loss Shape-Loss divides each pixel into edge pixels and internal pixels. For the infrared small target image I, the pixels in the loss function are set to be inside the frame, and the pixels outside the frame are removed by the bounding box mask. In order to generate the label value, a directional difference convolution kernel v' is used to generate the shape map V, and the thresholds λ and α are set, where λ>α. This makes the points in the shape map V greater than λ defined as edge pixels, and the points less than α are defined as internal pixels, thereby successfully generating the label value. For the predicted value Y', the corresponding edge shape map U is also required. x In the x-axis direction, the predicted values ​​y of the pixels on the left and right sides of the pixel e are obtained. i,j and y k,l , and then calculate the probability P(U e =1). The specific formula is as follows:

[0040] P(U e =1)=y i,j *y k,l +(1-y k,l )*(1-y i,j ), (6)

[0041] P(U e =0) = 1-P(U e =1). (7)

[0042] The same calculation method is applied to the y-axis direction. That is, the shape map H is generated using the differential convolution kernel h' in another direction, and then the same calculation and operation as the shape map V are performed. That is, formulas (6) and (7) are applicable to the edge shape map U corresponding to the shape map V. x , also applies to the edge shape graph U corresponding to the shape graph H y .

[0043] After obtaining the shape diagram, Figure 7 As shown, the Shape-Loss loss within the frame of the current infrared small target image is calculated, and the loss function is calculated as follows:

[0044] L shape =(I v＞λ log(P(U e =1))+(1-I v＜α )log(1-P(U e =0)) / (sum(V)) (8)

[0045] Among them, v refers to any label value in the shape graph V corresponding to the infrared small target image I; sum(V) means summing the label values ​​of the shape graph V.

[0046] To confirm this idea, the experiment visualized the relationship between the false positive rate of pixels and the threshold size on the infrared small target dataset NUDT-SIRST. Figure 8 The results show that when the thresholds λ and α are changed, the false positive rate of edge pixels decreases as the threshold λ increases, while the false positive rate of interior pixels decreases as the threshold α decreases. If the λ threshold is 0.2, only 2% of the pixels within the box are mistakenly classified as edges positive. Further increasing λ can further reduce this proportion, but excessively large thresholds reduce the number of supervised edges. Experiments have shown that the threshold is not sensitive to the final performance. Given a high proportion of positive edges, during training, it is safe to assume that all pixels with v ≥ λ are positive and only calculate the loss for these pixels. Other edges are discarded during the loss calculation. Similarly, when v ≤ α, all pixels are marked as interior pixels. Finally, the predicted values ​​are similarly mapped to interior and exterior pixels to obtain the shape-loss function for the interior of the box.

[0047] In the weakly supervised segmentation algorithm for infrared small targets, the combination of the framework based on the EMA algorithm and the weakly supervised loss function is explored. The EMA algorithm is often used to smooth time series data. By calculating the moving average, it suppresses noise and better captures trends. Borrowing the idea of ​​EMA, it is applied to the training framework of infrared small target segmentation. Equation (9) shows the change formula of the teacher model update, where f teacher (θ) and f student (θ) represents the specific parameters of the teacher model and the student model, respectively. β represents the learning iteration amplitude of the teacher and student models, which is in the range of (0, 1). During the iterative update of the algorithm, β is gradually approached from a very small value to 1, steadily improving the reliability of the pseudo mask and thus improving the learning stability of the model for small infrared targets:

[0048] f teacher (θ)=β*f student (θ)+(1-β)*f teacher ( (9)

[0049] Considering that box-supervised methods can generate high-quality masks without mask annotations, a self-training strategy can be adopted, using these high-quality masks as pseudo-labels to train fully supervised instance segmentation methods. Specifically, the BoxInst algorithm, which has been proven effective in weakly supervised image segmentation, is used to generate pseudo-instance masks on a given dataset, containing only box annotations. For each input image X, let B and M denote the bounding box mask and predicted mask, respectively. This paper proposes a simple yet effective weakly supervised segmentation algorithm for infrared small targets based on EAM, the pseudo-code of which is shown in Algorithm 1. The predicted mask M generated by the teacher model is corrected by the true bounding box mask B, and the corrected predicted mask M is assigned to the student model as the bounding box mask. The assigned instances are corrected by removing areas outside the bounding box. The student model is then trained using a dataset with pseudo-labels as instance masks. The teacher model is then updated using EAM and the student model, and iterative training produces higher-level masks. This training process, by leveraging high-quality pseudo-labels, compensates for the lack of original supervised data and improves the performance and generalization ability of the instance segmentation model.

[0050] Assuming the initial teacher model is not accurate enough, accurate coarse labels are needed to guide model training. As the teacher model becomes more accurate over time, the proportion of pseudo-mask losses can be increased. This involves updating the teacher model using the EMA technique, slowly adjusting it to a more accurate level through gradient coefficients. Furthermore, a weakly supervised loss function is introduced, using self-training to train a fully supervised instance segmentation method using high-quality masks as pseudo-labels. This method reduces the need for labeled data, making the model more suitable for infrared small target detection. Based on the EMA algorithm, we optimize the above infrared small target training method and design a weakly supervised training method for infrared small targets. In Algorithm 1, each change in the infrared small target image mask requires the algorithm to complete model training, resulting in an uncertain number of cycles, which is extremely time-consuming and unstable. This EMA-based weakly supervised training of infrared small targets requires the simultaneous initialization of two models. The algorithm runs both models simultaneously, subtly changing the mask model after each training cycle. This is an improved cyclic training method that significantly reduces training time. This EMA-based training method offers significant advantages for infrared small target weakly supervised algorithms. By running the two models in parallel, image masks of small infrared targets can be acquired and updated more flexibly, without being constrained by the uncertain number of training cycles and time overhead of traditional methods. This improved model design makes the algorithm more stable and efficient, and more adaptable to the characteristics of different small infrared targets. Furthermore, by fine-tuning the mask, model updates are faster, helping to improve the algorithm's convergence and performance. In practical applications, this weakly supervised training method has proven feasible and reliable for the rapid and accurate recognition of small infrared targets.

[0051] Models include

[0052]

[0053] By combining the EMA algorithm with a weakly supervised loss function, an innovative training framework is proposed to address the data labeling challenge in infrared small object segmentation, improving model performance and generalization. Subsequent experiments will verify the effectiveness of this framework and explore its potential applications and impact in the field of infrared small object segmentation. This research provides new insights and approaches for the advancement of infrared small object segmentation. Ablation experiments conduct an in-depth study of the proposed EMA framework and the weakly supervised loss function to evaluate their independent contributions and interactions in the task. In the validation experiments of the EMA framework, different components of the EMA algorithm are gradually introduced and the model performance is compared. The results show that the introduction of the EMA algorithm significantly improves the robustness and training stability of the model, effectively suppresses noise interference during training, and enables the model to better adapt to the characteristics of infrared small objects. Secondly, in the validation experiments of the weakly supervised loss function, the impact of different forms of weak supervision signals on model performance is analyzed. By comparing models trained with different pseudo-labels, the importance of the weakly supervised loss function in improving model performance is demonstrated. Especially when labeled data is limited, the weakly supervised loss function provides more effective information to the model, which helps improve the model's generalization ability. Finally, we conducted joint experiments with the EMA framework and a weakly supervised loss function to verify their synergistic effect. The experimental results show that the EMA framework and the weakly supervised loss function complement each other and work together to further improve the model's performance on the infrared small object segmentation task. These ablation experiments provide strong support for the rationality and effectiveness of the proposed method, and offer a practical solution for research and application in the field of infrared small object segmentation.

[0054] Experimental design and result analysis

[0055] Dataset

[0056] The existing dataset needs to be converted into a box-annotated dataset to simulate a box mask dataset. Specifically, in this process, the same infrared small target image is used to convert the infrared small target annotation mask into a matrix with the same length and width to better simulate the box-annotated dataset. Figure 9 The process and results of this conversion using manual annotation to simulate coarse labels are shown.

[0057] This image mask conversion is designed to simulate weakly supervised labels, providing a more challenging dataset for the study of small infrared targets. In this process, the weakly supervised labels are simulated by converting the infrared small target annotation mask into a corresponding matrix form. The main purpose is to provide the model with approximate information about the target location in the absence of detailed annotation information. Through such weakly supervised labels, it is possible to more closely resemble the high labeling costs in real-world application scenarios, while also challenging the model's robustness when dealing with data that lacks precise annotations.

[0058] When faced with large-scale infrared small target image datasets, manually annotating the precise location of each small target is a tedious and expensive task. By introducing such weakly supervised labels, we can better address the challenges of data acquisition and annotation, providing a more practical and feasible solution for training infrared small target segmentation models. This data processing strategy aims to promote researchers' in-depth exploration of the infrared small target segmentation problem and provide a more operational data foundation for the development of related fields.

[0059] Experimental parameter settings

[0060] Experiments were conducted on a Linux computer equipped with a 2.50GHz CPU, 11GB of RAM, and a GeForce GTX 2080ti GPU. Models were implemented in Python and PyTorch. Training was performed on ID-Unet using the Intersection over Union (IoU) loss function and optimized using the Adgrad method. The parameters λ and α of the Shape-Loss within the bounds were 0.2 and 0.03, respectively. The parameter β of the EMA teacher model algorithm was 1 / N, where N is the number of iterations. The model weights and biases were initialized using the Xavier method, and all models were trained from scratch.

[0061] Comparative experiment

[0062] Due to the differences between natural images and images of small infrared targets, most weakly supervised image segmentation methods are not optimized for the characteristics of small infrared targets. Table 1 shows the comparative experimental results of our method and other weakly supervised image segmentation algorithms on a small infrared target dataset. Full supervision refers to the results using manually annotated ground truth masks, while bounding box masks refer to the results using bounding box masks directly without weak supervision. The experimental results show that our method outperforms other methods on all three datasets.

[0063] The proposed method was compared with weakly supervised natural image segmentation methods on the NUDT-SIRST, NUAA-SIRST, and IRSTD-1k datasets. The proposed method employed the infrared weakly supervised loss function PS-Loss, which includes the out-of-box loss Project Loss and the in-box loss Shape-Loss, and introduced an EMA-based dual-model iterative algorithm, ultimately achieving satisfactory experimental results. Compared with the current natural image segmentation algorithm BoxTeacher, the proposed method achieved improvements of 2.69%, 4.43%, and 1.67% in IoU, nIoU, and F1, respectively, on the NUDT-SIRST dataset; 0.99%, 0.62%, and 0.65% in IoU, nIoU, and F1, respectively, on the NUAA-SIRST dataset; and 2.05% and 1.48% in IoU and F1, respectively, on the IRSTD-1k dataset, while reducing nIoU by 0.32%. Our method optimizes and improves performance on three infrared small target datasets, achieving higher image segmentation metrics, such as Intersection of Union (IoU) and F1 scores. By combining out-of-frame and in-frame losses, we successfully reduce information loss for small infrared targets, thereby improving the accuracy and robustness of segmentation results. Overall, our method demonstrates significant advantages in addressing the challenges of infrared small target segmentation, providing an effective solution to addressing this field.

[0064] Table 1 Experimental comparison results with other weakly supervised image segmentation methods

[0065]

[0066]

[0067]

[0068] The introduction of an out-of-box loss effectively confines small infrared targets within a predefined box, preventing the loss of small target information during segmentation. This ensures the model's focus is on the region of interest, thereby improving segmentation accuracy. Experimental results show that the use of the out-of-box loss reduces the model's influence on the out-of-box region, effectively reducing the model's sensitivity to background and improving overall stability. The introduction of the in-box loss enhances the modeling and learning of target shape by adjusting the shape of the small target within the box. Shape information is crucial to the success of infrared small target segmentation. Experimental results show that the in-box loss effectively improves the model's accurate understanding of the small target's shape, making the segmentation results more realistic and further enhancing model performance. BoxInf leverages the advantages of both the out-of-box and in-box losses, achieving accurate segmentation of small infrared targets by comprehensively considering target position and shape information. Compared to other methods, it achieves significant advantages in both accuracy and robustness, demonstrating superior segmentation performance. Due to its improved loss function, the BoxInf method is universally applicable to different detection models and is suitable for a variety of infrared small target segmentation models. Experimental results demonstrate that our method generalizes well across different scenarios, background noise, and small target variations, demonstrating its potential for practical applications. In summary, the BoxInf method performs well in the infrared small target segmentation task and has significant advantages in improving segmentation accuracy and robustness.

[0069] Compared with other methods, the experimental results of our method are shown in the table. The most advanced methods in the field of image segmentation are usually difficult to train and segment effectively in the field of infrared small targets. Our method experimentally compares these methods on three datasets, and the experimental results are shown in the table. The results show that our method works best on weakly supervised infrared small target images. The most advanced algorithms usually have difficulty in effectively segmenting and subdividing infrared small targets that change in complex scenes encountered in the NUDT-SIRST dataset. In contrast, BoxInf shows greater robustness to these changes. In addition, compared with the most advanced algorithms, BoxInf can more accurately depict the target shape. Our innovative backbone design effectively eliminates noise by fusing deep information and adapts to various clutter backgrounds, target shapes and target size challenges, thereby achieving excellent performance. As Figure 10 As shown in Figure 3, we applied the BoxInf model and compared it with the state-of-the-art models using false positive rate and true positive rate. By plotting ROC curves for different thresholds, we analyzed and concluded that the BoxInf model surpassed the existing best models.

[0070] Ablation experiments

[0071] The present invention uses ablation experiments to verify the effectiveness of the infrared weak supervision method and the influence of some modules, including the influence of two IoU loss functions on the method, the combined ablation of weak supervision loss functions, and the influence of the EMA-based iterative algorithm.

[0072] Table 2 Ablation experiments with different loss function combinations

[0073]

[0074] (a) Impact of IoU loss function. In the ablation experiment, we focused on the impact of the IoU loss function on the performance of our method. In Table 2, separate experiments were conducted using two different loss functions to gain a deeper understanding of their performance under experimental conditions. The experimental results show that when the labels are inaccurate for the method, using the SoftIoU-Loss function can achieve better results and alleviate the negative impact of inaccurate labels on model performance. In contrast, when the labels are accurate for the method, using the DiceIoU-Loss function performs better. In this case, DiceIoU-Loss can better capture the similarity of the target area, thereby improving the accuracy of segmentation. This series of experiments proves that choosing the appropriate loss function in different scenarios is crucial to the performance of infrared weakly supervised methods. The selected loss function directly affects the model's ability to adapt to label uncertainty and accuracy. Therefore, in practical applications, choosing the most appropriate loss function based on the characteristics of the specific scenario is an important factor in ensuring the superior performance of the model.

[0075] Table 3 Ablation experiments on the module composition of the BoxInf model

[0076]

[0077] (b) Contribution of each module. Contribution of each loss term. The final loss function consists of two loss functions, and Table 3 shows the contribution of each loss term. Even if only the first projection term Project-Loss is used, satisfactory performance can be achieved. This is much higher than using only the bounding box mask. By further using the proposed edge loss Shape-Loss, the weakly supervised segmentation capability of small infrared targets can be further improved. Finally, the two loss functions are integrated into the EMA-based iterative algorithm to further improve the model performance. This series of experimental results shows that the reasonable combination and adjustment of different loss terms have a significant impact on the improvement of overall performance. The weakly supervised method using all modules achieves 1.25%, 1.22%, and 0.78% improvements in IoU, nIoU, and F1, respectively, compared to the method without EMA. It also achieves 3.42%, 3.78%, and 2.14% improvements in IoU, nIoU, and F1, respectively, compared to the weakly supervised method using only the Project loss function. Furthermore, it achieves 23.67%, 18.45%, and 6.74% improvements in IoU, nIoU, and F1, respectively, compared to the method using only bounding box masks for full supervision. These experimental results demonstrate the importance of each module in our method.

[0078] Table 4 Similarity between preprocessed data and true masks

[0079]

[0080]

[0081] (c) The impact of different data preprocessing methods. During the dataset preprocessing stage, the dataset is processed through a series of operations, and the processed dataset is directly used for the weakly supervised task of infrared small target segmentation. In this task, the performance of the method is evaluated using indicators such as hit rate, IoU, and nIoU. The higher the values ​​of these indicators, the better the effect on subsequent weakly supervised training. Experiments were conducted using multiple methods, including the GrabCut method, which uses the GrabCut algorithm to segment small targets within the box. In addition, the box method was introduced, which sets the pixel values ​​within the box annotation to the target pixel. GrabCut is combined with the Box method. This method first uses the GrabCut algorithm to segment small targets within the box, and then uses the box method to generate masks for undetected small targets. Through comparative experiments of these methods, the goal is to find the data processing and segmentation methods that are most effective in weakly supervised tasks to optimize the performance of infrared small target segmentation.

[0082] Table 5 Final training results of different preprocessing methods

[0083]

[0084] (d) Experimental results of data preprocessing using different methods. As Table 5 shows, the hit rate of using GrabCut alone is far inferior to that of using bounding box masks and the combined method. Due to its limitations, the algorithm cannot segment all small objects. The experimental results of this method are clearly poor. In Table 5, we experimentally compare the box method trained with bounding box masks for full supervision with the BoxGrabCut method that combines the advantages of both methods. Experimental results show that the combined method outperforms the box method trained with bounding box masks for full supervision on all three datasets. On the NUDT-SIRST dataset, IoU, nIoU, and F1 are improved by 8.7%, 3.62%, and 6.71%, respectively; on the NUAA-SIRST dataset, IoU and F1 are improved by 0.55% and 0.38%, respectively, while nIoU is reduced by 0.5%; and on the IRSTD-1K dataset, IoU, nIoU, and F1 are improved by 1.82%, 4.25%, and 1.45%, respectively. Experimental results demonstrate the effectiveness of the combined method in image preprocessing, which can effectively improve the accuracy of subsequent weakly supervised methods.

[0085] Table 6 Final effects of different β pair methods in EMA iterative algorithm

[0086]

[0087] (e) The impact of the weak supervision parameter λ on the experimental results. Table 6 shows the impact of different λ parameters in the weak supervision method of the present invention. In the experiment, we set the maximum value of the incremental λ to 1, and experimented with the impact of different λ values ​​on the total number of training iterations N and the performance indicator IoU. Specifically, we selected experimental settings when λ was 1 / N, 2 / N, and 3 / N, and observed their impact on IoU. The experimental results show that when λ is 1 / N, IoU is only improved by 0.2% compared to when λ is 2 / N, and by 0.35% compared to when λ is 3 / N. This shows that the impact of the λ parameter on the method is not particularly significant. It is worth noting that the choice of λ may need to be fine-tuned in other experiments to better adapt to different infrared small target datasets. This flexibility makes our method more versatile and adaptable, and can achieve better performance in dealing with different data scenarios.

[0088] Table 7 Ablation experiments on loss function combinations

[0089]

[0090]

[0091] (f) Impact of labeling errors. During the coarse labeling process, some errors may be introduced due to various reasons. In order to test the robustness of our method to such errors, we conducted a series of simulation experiments, as shown in Table 7. Specifically, we simulated the errors that may occur in the labeling process by introducing an offset. The offset here is used to indicate the degree of error, where +1 means an extra circle of pixel values, +2 means an extra circle of pixel values, and +3 means an extra circle of pixel values. Through such simulations, we can evaluate the robustness of our method in the face of labeling errors, providing more comprehensive and reliable results for our research.

[0092] Experiments have shown that the infrared small target segmentation task is highly sensitive to labeling accuracy. A small infrared target may consist of only a few pixels. Each time the pixel values ​​are artificially expanded, the detection performance of the weakly supervised infrared target algorithm is significantly reduced, demonstrating that labeling accuracy significantly impacts the final segmentation efficiency. Experiments also demonstrate that a decrease in manual labeling accuracy due to increased data volume can also affect final model performance, further highlighting the importance of weak supervision.

[0093] Table 8 Ablation experiment of PSLoss loss function weighted value

[0094]

[0095] (g) The influence of the weighted value of the loss function. The final loss function consists of two loss functions, divided into out-of-box loss and in-box loss, in which the results of different weighted values ​​are tested. Table 8 shows the impact of different in-box loss weighted values ​​on the experimental results. The weighted value of the loss function plays a vital role in the training process, affecting the convergence speed and final performance of the model. By adjusting the weighted value of the in-box loss, the importance of the model in capturing the internal details and overall shape of the target can be balanced. The experimental results show that a suitable weighted value of the in-box loss can better promote the model's learning of the internal details of the target, but an excessively high weighted value may cause the model to focus too much on the internal details of the target and ignore the overall shape, thereby affecting the accuracy of the segmentation. Therefore, in practical applications, it is necessary to reasonably select the weighted value of the loss function according to the characteristics of the specific task and dataset to achieve the best segmentation effect.

[0096] In summary, the present invention aims to address the challenges of small size, sparse number, and difficult labeling in infrared small target images. The comprehensive weakly supervised learning framework designed in this invention simulates fine labeling by introducing image mask transformations, thereby reducing labeling complexity. Experimental validation using three different datasets demonstrates the effectiveness of BoxInf. Experimental results show that BoxInf achieves significant performance improvements in the infrared small target segmentation task. By employing a training strategy that combines coarse and fine labels, it achieves excellent performance in both target detection and segmentation accuracy. In particular, it demonstrates superior robustness compared to traditional methods in complex scenes and various background clutter. Comparison of experimental results across different datasets verifies the algorithm's generalization and adaptability. This innovative application of weakly supervised methods to the infrared small target segmentation task offers a new approach to solving this problem. By designing a weakly supervised method, the cost of labeling is reduced, increasing the number of infrared small target datasets, addressing the current bottleneck in the development of infrared small target tasks and making infrared small target segmentation technology more practical. The present method not only demonstrates significant advantages in labeling costs but also provides an innovative solution for infrared small target detection while maintaining segmentation quality. The introduction of this new method will promote the continuous development of the field of infrared small target segmentation and provide more feasible and economical solutions for practical applications.

[0097] This paper proposes a simple, high-performance, and end-to-end weakly supervised segmentation algorithm for infrared small target bounding boxes. This method is based on the ID-UNet model and proposes a high-performance weakly supervised method that can achieve mask-level instance segmentation of infrared small targets using only the bounding box annotations used for training. Only for newly acquired infrared small target data, the data needs to be annotated with labeled boxes, which greatly reduces the annotation time. Considering the huge cost of labeled instance segmentation, weakly supervised instance segmentation using image-level labels or bounding boxes has received widespread attention. However, we are also aware of the challenges of the algorithm in handling some special cases, such as the inability to effectively and accurately obtain the shape edges of small targets with strange shapes. Future work can further improve the performance of the algorithm by introducing more contextual information or adopting an attention mechanism. Overall, our research provides a new solution to the problem of scarce datasets in the infrared small target segmentation task, providing beneficial assistance for the practical application of segmentation algorithms.

Claims

1. A weakly supervised infrared small target segmentation method, characterized in that: The input is only the infrared small target image and the infrared small target bounding box data. The final output result is the segmentation mask of the small target in the infrared small target image. The method includes model pre-training, model loading, teacher model generation of pseudo mask, student model training and dual-model EMA. The teacher model generation of pseudo mask, student model training and dual-model EMA need to be repeated for multiple iterations. The specific steps include the following: Step 1: Use the infrared small target image and the corresponding bounding box mask to pre-train the model in the PS-Loss loss function; Step 2: Copy the parameters of the pre-trained model to the teacher model and the student model respectively; fix the parameters of the teacher model for use in subsequent steps; Step 3: Generate a pseudo mask from the infrared small target image through the teacher model; Step 4: Compare the pseudo mask with the bounding box mask and remove the pseudo mask outside the bounding box. This is because the infrared small target has been pre-set to be within the bounding box. Step 5: Use the bounding box mask from step 1 and the pseudo mask processed in step 3 as guidance labels for the loss function to control the student model training and obtain the output of the student model. In the process of training the student model through the loss function, the bounding box mask uses the PS-Loss loss function, while the pseudo mask uses the general image segmentation loss function DiceIoU-Loss loss function. Step 6: Input the student model and the teacher model into the EMA algorithm to update the teacher model, then return to step 4 and repeat the iterative training multiple times. The final student model is used as the infrared small target segmentation model, and the predicted mask output by the final student model is used as the final output for evaluation; The loss function PS-Loss of weak supervision of infrared small targets consists of the projection loss Project-Loss outside the frame and the shape loss Shape-Loss inside the frame. The weak supervision loss L ps Defined as: THE ps =L project +λL shape (1) Among them L project is the out-of-frame projection loss, with a weight of 1; L shape It is the in-box shape loss, and the weight value defaults to λ.

2. The infrared small target segmentation method based on weak supervision according to claim 1 is characterized in that: The loss function formula of the out-of-frame projection loss Project-Loss is as follows: Proj x (m)=max y (m)=m x , (2) Proj y (b)=max x (b)=b x , (3) L Project =L DiceIoU (m x ,b x )+L DiceIoU (m y ,b y ), (4) The loss function is to minimize the difference between the bounding box mask and the predicted mask, and the corresponding projection m on the x-axis and y-axis is obtained by calculating the maximum value of the bounding box mask on the y-axis and x-axis x and m y , predict the corresponding projection b of the mask on the x-axis and y-axis x and b y , max y (m) represents the mapping of the bounding box mask on the y-axis, and max x (b) represents the mapping of the predicted mask on the x-axis; L DiceIoU Represents the DiceIoU-Loss function.

3. The infrared small target segmentation method based on weak supervision according to claim 2 is characterized in that: Shape-Loss uses a directional difference convolution kernel to extract the shape features of small infrared targets and pre-sets labels based on the shape features to generate prior labels. For infrared small target image I∈R w*h , use the x-axis and y-axis direction differential convolution kernels v' and h' to extract the shape feature map U of the infrared small target in two directions x and U y , and then use the square sum of the two variables to generate the shape; by setting the threshold, the larger and smaller values ​​are set as edge pixels and internal pixels respectively, and the edge shape map U∈R is generated by edge pixels and internal pixels w*h , guiding the model to perform training optimization; the formula for obtaining the edge shape graph is as follows: The in-frame shape loss Shape-Loss divides each pixel into edge pixels and internal pixels. For the infrared small target image I, the pixels in the loss function are set to be inside the frame, while the pixels outside the frame are removed by the bounding box mask. In order to generate the label value, a directional difference convolution kernel v' is used to generate the shape map V, and the thresholds λ and α are set, where λ>α. This makes the points in the shape map V greater than λ defined as edge pixels, and the points less than α are defined as internal pixels, thereby successfully generating the label value. For the predicted value Y', the corresponding edge shape map U is also required. x ; In the x-axis direction, the predicted value y of the pixel points on the left and right sides of the pixel point e is obtained i,j and y k,l , and then calculate the probability P(U e =1); the formula is as follows: P(U e =1)=and i,j *and k,l +(1-and k,l )*(1-and i,j ) (6) P(U e =0)=1-P(U e =1). (7) The same calculation method is applied to the y-axis direction; that is, the shape map H is generated using the other direction differential convolution kernel h', and then the same calculation and operation as the shape map V are performed; that is, formulas (6) and (7) are applicable to the edge shape map U corresponding to the shape map V. x , also applies to the edge shape graph U corresponding to the shape graph H y ; After obtaining the shape map, the Shape-Loss loss within the frame of the current infrared small target map is calculated. The loss function is calculated as follows: L shape =(I v＞λ log(P(U e =1))+(1-I v＜α )log(1-P(U e =0)) / (sum(V)) (8) Among them, v refers to any label value in the shape graph V corresponding to the infrared small target image I; sum(V) means summing the label values ​​of the shape graph V.

4. The infrared small target segmentation method based on weak supervision according to claim 3 is characterized in that: This method combines the framework based on the EMA algorithm with the weakly supervised loss function; Equation (9) shows the change formula for updating the teacher model, where f teacher (θ) and f student (θ) represents the specific parameters of the teacher model and the student model, and β represents the learning iteration amplitude of the teacher and student models, which is in the range of (0,1]. During the iterative update, β is gradually approached from a very small value to 1, steadily improving the reliability of the pseudo mask and thus improving the learning stability of the model for small infrared targets: f teacher (θ)=β*f student (θ)+(1-β)*f teacher (i).(9) These high-quality masks are used as pseudo labels to train a fully supervised instance segmentation method. Specifically, the BoxInst algorithm, which has been proven to be effective in the field of weakly supervised image segmentation, is used to generate pseudo instance masks on a given dataset, which only contain box annotations. For each input image X, let B and M represent the bounding box mask and the predicted mask, respectively. The predicted mask M generated by the teacher model is corrected by the true bounding box mask B, and the corrected predicted mask M is assigned to the student model as the bounding box mask. The area outside the bounding box is deleted to correct the assigned instance, and then the pseudo-label is used as the dataset of the instance mask to train the student model. The teacher model is then updated through EMA and the student model, and iterative training produces a higher level mask.

Images