Semi-supervised object detection method based on feature repository and contrastive learning

Abstract

A semi-supervised target detection method based on feature repository and contrast learning is provided, the feature repository is used to store high-quality feature vectors, and the feature vectors of unlabeled data are constrained by contrast learning, so that additional supervision information of labeled data on unlabeled data can be established, the training of the model is regulated, the direction of training is optimized, and the performance of the model is improved, which has practical significance and good application prospect.

Description

Technical Field

[0001] This invention relates to the fields of image processing and computer vision, and in particular to a target detection method under semi-supervised learning. Background Technology

[0002] Object detection is a popular area in computer vision and digital image processing. It enables the automatic detection of object categories and locations, reducing the consumption of manpower and capital, and has significant practical implications. Compared with more basic image classification tasks, object detection adds a regression task. That is, it not only needs to use algorithms to determine whether an object exists in the image, but also to mark its location in the image and perform regression prediction on the marked bounding boxes on the image.

[0003] In recent years, object detection algorithms have developed rapidly due to the widespread application of deep learning. However, typical fully supervised object detection models are trained on large datasets with precise manual annotations. These methods require each training image to have accurate and sufficient high-quality annotations. Often, an image contains multiple objects, each potentially belonging to a different category, all of which require manual annotation. Furthermore, some objects are difficult to identify with the naked eye due to their small size, environmental influences, or image distortion, further increasing the time and effort required for dataset annotation. Therefore, to reduce the significant human cost of annotation, how to fully utilize small-sample labeled datasets has become a major research focus.

[0004] Based on this, semi-supervised learning attempts to enhance model performance by utilizing more readily available unlabeled data on small samples of labeled data. Currently, most semi-supervised object detection methods are limited by issues such as low pseudo-label accuracy, resulting in a significant gap in performance compared to fully supervised object detection methods. Therefore, further research on semi-supervised object detection methods is of great significance. Summary of the Invention

[0005] Current semi-supervised object detection methods, both domestically and internationally, fail to fully leverage supervisory information. Existing methods typically employ consistency loss mechanisms. For example, Chinese patent CN112926673A discloses a semi-supervised object detection method based on consistency constraints. This method calculates consistency loss on prediction results for different image enhancements to optimize the model, but it fails to further model the relationship between labeled and unlabeled data. Another example is Chinese patent CN114399683A, which discloses an end-to-end semi-supervised object detection method based on an improved YOLOv5. This method iteratively generates pseudo-labels for new model training, but it is significantly affected by pseudo-label noise, requiring the introduction of additional information to supervise and guide pseudo-label generation.

[0006] To address this, this paper designs a semi-supervised object detection method based on a feature repository and contrastive learning. The feature repository is used to store high-quality feature vectors, and contrastive learning is used to constrain the feature vectors of unlabeled data. This can establish additional supervision information of labeled data on unlabeled data, standardize model training, optimize the training direction, and improve model performance. This method has practical significance and good application prospects.

[0007] The specific contents of this invention are as follows:

[0008] One approach based on the following steps:

[0009] The first step is to obtain the semi-supervised object detection dataset {D}. l D u}, of which the labeled portion of the dataset is Unlabeled dataset is N l N u These represent the number of labeled images and the number of unlabeled images, respectively. For image The tag information. i Let b be the truth class label for each bounding box in the i-th image. i The true value of each bounding box in the i-th image is assigned to the coordinate label of the bounding box.

[0010] The second step is to build an object detection model, which includes a feature extraction module. Region generation module Output prediction module Where f is the feature vector of the input prediction module, and b is the selected bounding box location information; each module is composed of several layers of convolutional neural network or fully connected layers stacked together. For input x i Region generation module The output region generation result is Output prediction module The output prediction result is in For the predicted class vector, These are the predicted values ​​for each calibration frame.

[0011] The third step is to establish a feature repository. This repository is used to store feature vectors for each category. Its size is C×T×D, where C is the number of categories in the dataset, T is the number of vectors for a specific category stored in the repository, and D is the dimension of the stored feature vectors. The repository uses a First-In-First-Out (FIFO) queue storage model. If the total number of currently stored vectors exceeds the size T, the vectors added to the repository first will be removed, and the newest feature vectors will be added.

[0012] The fourth step is to establish a feature mapping module. Map the feature f corresponding to the selected region b to a lower-dimensional feature vector f. proj ;

[0013] Step 5, for labeled data D l Image data in Obtain the truth value calibration box annotation Region corresponding to mapping feature vector The corresponding categories for each region are: According to category Map the feature vectors of each region Store in feature repository middle

[0014] Step 6, for unlabeled data D u Image data in First, use the region generation module to obtain the bounding box of the region of interest. Reliable feature vectors are selected, and then the corresponding mapped feature vectors for each bounding box region that meets the criteria are obtained. The corresponding categories for each region are: According to category Map the feature vectors of each region Store in feature repository middle.

[0015] Step 7: Construct the contrastive learning loss function. For unlabeled data... All prediction results If the prediction category is calibration frame features If a sample is taken, its corresponding contrastive loss function is in the form of: This represents the number of sampling calibration boxes for the current image. Where P... v and N v These are sets of positive and negative sample feature vectors, respectively. The feature vectors in these sets come from a feature repository, P. v The positive sample set consists of v u The vector composition of N belonging to the same category c v It consists of vectors from all other categories; sim(·,·) is a vector similarity metric.

[0016] Step 8, with labeled data D l Use the usual classification loss. and regression loss The total loss function for labeled data is used as the loss function during training. In unlabeled data D u The above uses self-training loss. or consistency loss Optimize the model, denoted as The self-training method uses a teacher-student model, with the teacher model generating pseudo-labels for the student model to train. The consistency method uses different transformations of an image as input and calculates the difference in predictions under different transformations. The total loss function for unlabeled data is...

[0017] Step 9: Optimize the loss function using stochastic gradient descent, train the model, and repeat steps 5-8, iterating multiple times until the model's prediction results are stable and optimal.

[0018] The specific functions and steps of each module in the second step include:

[0019] I. Feature Extraction Module Image feature extraction. This module first preprocesses the training images, and then uses a common feature extraction network to extract features from the training images.Typical feature extraction networks include VGG (https: / / arxiv.org / pdf / 1409.1556, Very Deep Convolutional Networks, 2014), ResNet (https: / / arxiv.org / pdf / 1512.03385.pdf, Deep Residual Learning for ImageRecognition, 2015), MobileNet (https: / / arxiv.org / pdf / 1704.04861, MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications, 2017), RetinaNet (https: / / arxiv.org / pdf / 1708.02002, Focal Loss for Dense ObjectDetection, 2018) or EfficientNet (https: / / arxiv.org / pdf / 1905.11946, EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks). Furthermore, feature extraction networks are combined with feature processing networks to further extract and optimize image features. Typical feature processing networks include BAM (https: / / arxiv.org / pdf / 1807.06514, BAM: Bottleneck Attention Module, 2018), CBAM, SPP (https: / / arxiv.org / pdf / 1406.4729, Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition, 2014), FPN (https: / / arxiv.org / pdf / 1612.03144, Feature Pyramid Networks for Object Detection, 2016) and / or NAS-FPN (https: / / arxiv.org / pdf / 1904.07392, NAS-FPN: Learning Scalable Feature Pyramid Architecture for Object Detection, 2019).

[0020] II. Region Generation Module Generate candidate regions, with the input being the image x. i The generated candidate regions (label boxes) The image features (f) and the image data (f) serve as inputs to the output prediction module. This module is either the dense bounding box generator in a one-stage object detection method or the Region Proposal Network (RPN) in a two-stage object detection method to generate candidate regions.

[0021] III. Output Prediction Module Predict the type c and location b of the target. This module first performs RoI Pooling / Align on the features within the candidate region, which is either a one-stage dense prediction YOLO network (https: / / arxiv.org / pdf / 1506.02640, You Only Look Once: Unified, Real-Time Object Detection, 2015) or a two-stage refined prediction Faster R-CNN network (https: / / arxiv.org / pdf / 1506.01497, Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks, 2015) or an anchorless CenterNet network;

[0022] In the third step, to avoid excessive GPU memory usage, the feature dimension D is usually set to 128 or 256.

[0023] In the fourth step, the feature mapping module The MLP (Multilayer Perceptron) is adopted, and its structure is: fully connected layer -> ReLU activation layer -> batch normalization -> fully connected layer.

[0024] In the fifth and sixth steps, the feature vectors are stored in the feature repository. If the corresponding category in the library is full, the earliest added feature vector will be removed according to the first-in-first-out principle, and the existing new feature vector will be stored in.

[0025] The reliable feature vector selection in step six is ​​determined by the category prediction scores of each region of interest candidate box. If the highest category prediction score of a region of interest candidate box is greater than a threshold θ, then the feature vector corresponding to that candidate box is considered a reliable feature vector.

[0026] In the seventh step, the similarity metric sim(·,·) has multiple implementations, which can be used for any two different feature vectors x. iand x j Using vector dot product MLP layer mapping Calculate the cosine distance between vectors Gaussian method This is used to measure the similarity between two samples. The role of the contrastive learning loss function is to bring feature samples of the same class closer together in the high-dimensional feature space, while pulling samples of different classes further apart. Combined with a feature repository, it can effectively establish the guidance of labeled information for unlabeled data.

[0027] In the ninth step, the criterion for terminating the iteration must be that the final loss of the model tends to stabilize in the last few iterations without a significant decrease, rather than the loss stabilizing in a single iteration.

[0028] Compared with the prior art, the present invention has the following innovative features:

[0029] 1. Other techniques using pseudo-labels rely solely on manually setting thresholds to filter reliable pseudo-labels. This results in poor-quality pseudo-labels with significant noise from misclassification or localization errors, hindering model training. This method extracts additional information from labeled data, providing reliable training direction for unlabeled data. It guides the model to generate high-quality pseudo-labels, and contrastive learning ensures this is achieved, ultimately resulting in better model prediction performance.

[0030] 2. Other techniques using consistency loss fail to establish an informational link between unlabeled and labeled data, lacking a channel for information communication across subsets. Feature repositories can establish this link and use contrastive learning for standardization and guidance, resulting in better model prediction performance.

[0031] 3. Experiments showed that, with other factors kept constant, the average accuracy (mAP) of the model without feature repository and with contrastive learning was 21%, while the mAP of the model with the latter method reached 25%, demonstrating the effectiveness of the method (this model is only used for testing the effectiveness of the method, and its numerical results do not represent the actual performance of the model).

[0032] The method execution flow provided in the embodiments of this application can be run on devices such as personal computers, servers, embedded computing devices, and cloud computing platforms. Attached Figure Description

[0033] To more clearly illustrate the technical solutions in the embodiments of this application 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 some embodiments recorded in this application. For those skilled in the art, other drawings can be obtained based on these drawings.

[0034] Figure 1 This is a flowchart of a semi-supervised object detection method based on feature repository and contrastive learning according to the present invention.

[0035] Figure 2 This is an example of DOTA dataset annotation used in one embodiment of the present invention.

[0036] Figure 3 This is a schematic diagram of the labeled image training process in one embodiment of the present invention, using the ResNet+FPN feature extraction module, the RPN region generation module, and the Faster RCNN output prediction module.

[0037] Figure 4 This is a schematic diagram of the unlabeled image training process in one embodiment of the present invention, using the ResNet+FPN feature extraction module, the RPN region generation module, and the Faster RCNN output prediction module. Detailed Implementation

[0038] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this application. All other embodiments obtained by those skilled in the art based on the embodiments of this application without creative effort are within the scope of protection of this application.

[0039] According to the embodiments of this application, the data used comes from the DOTA1.5 remote sensing image dataset (https: / / captain-whu.github.io / DOTA / dataset.html). Remote sensing images, compared to ordinary images in daily life, are characterized by small targets, densely packed objects, and arbitrary orientations, presenting greater challenges. The DOTA1.5 dataset itself is a fully supervised dataset. To construct a semi-supervised dataset for experimentation, some labels were discarded, resulting in labeled data accounting for ρ% of all training data. This value can range from 1% to 50%. In the experiment, it was set to 10%.

[0040] 1. Obtain the DOTA1.5 remote sensing image dataset. Select 15,000 images as training data, of which 1,500 images retain their annotation information as the labeled dataset, and 13,500 images are used as the unlabeled dataset. Figure 3 The image set is input.

[0041] 2. Establish a feature repository with a size C×T×D of 16*2000*256.

[0042] 3. For example Figure 3The diagram shows a labeled image training process. A specific example is: ResNet+FPN is used as the feature extraction module to extract image x. i Image features f i ResNet is responsible for backbone feature extraction, while FPN performs further optimization by performing multi-scale fusion detection on the extracted features to improve the accuracy of small object detection; the output feature map f i The data is fed into the Faster R-CNN output prediction module to predict and classify bounding boxes. The classification loss is calculated by comparing the predicted results with the ground truth labels. and regression loss Since the labeled images have accurate ground truth labels, the ground truth bounding boxes b in the dataset can be used directly. GT (without needing to use the PRN region generation module to predict the calibration box) Then, the corresponding feature vectors are extracted by the feature mapping module. Store it in the feature repository.

[0043] 4. For example Figure 4 The diagram illustrates the unlabeled image training process. A specific example is as follows: ResNet+FPN is used as the feature extraction module to extract image features. ResNet is responsible for extracting the backbone features, while FPN performs further optimization, performing multi-scale fusion detection on the extracted features to improve the accuracy of small object detection; the output feature map f is then processed. i The data is fed into the Faster R-CNN output prediction module for bounding box prediction and classification. The output prediction results are then used for training with the appropriate loss function based on the method employed. Since unlabeled images lack ground truth label information, the bounding box features generated by the feature repository and the RPN region generation module are used. Calculate the contrast loss (denoted as ). On the other hand, high-confidence prediction samples can also be added to the feature repository to expand their feature expression range.

[0044] 5. Combine labeled and unlabeled images into a batch and feed them into the model for training until the model's loss remains stable without significant decrease over several iterations. Use this stable model as the final trained model.

[0045] Although preferred embodiments of this application have been described, those skilled in the art, upon learning the basic inventive concept, can make other changes and modifications to these embodiments. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments as well as all changes and modifications falling within the scope of this application. Clearly, those skilled in the art can make various alterations and variations to this application without departing from its spirit and scope. Thus, if such modifications and variations fall within the scope of the claims of this application and their equivalents, this application also intends to include such modifications and variations.

Claims

1. A semi-supervised object detection method based on a feature repository and contrastive learning, characterized in that, Includes the following steps: The first step is to obtain a semi-supervised object detection dataset. The labeled portion of the dataset is The unlabeled dataset is , , These represent the number of labeled images and the number of unlabeled images, respectively. For image Tag information, Let each bounding box in the i-th image be a truth class label. The true value of each bounding box in the i-th image is assigned to the coordinate label of the bounding box. The second step is to build an object detection model, which includes a feature extraction module. Region generation module Output prediction module Where f is the feature vector of the input prediction module, and b is the selected bounding box location information; each module is composed of several layers of convolutional neural network layers or fully connected layers stacked together; for the input Region generation module The output region generation result is Output prediction module The output prediction result is ,in For the predicted class vector, These are the predicted values ​​for each calibration box; The third step is to establish a feature repository. Used to store the mapping feature vectors of various categories. The feature repository size is , For the number of categories in the dataset, Store the number of vectors for a certain category in the feature repository. The dimension for storing feature vectors; The fourth step is to establish a feature mapping module. The mapping module maps the features f of the region corresponding to the bounding box b to a lower-dimensional feature vector. ; Step 5, for labeled data Image data in Obtain the truth value calibration box annotation. Region corresponding to mapping feature vector The corresponding categories for each region are: According to category Map the feature vectors of each region Store in feature repository middle; Step 6, for unlabeled data Image data in First, the region generation module is used to obtain the bounding box of the region of interest. Reliable feature vectors are selected, and then the corresponding mapped feature vectors for each bounding box region that meets the conditions are obtained. The corresponding categories for each region are: According to category Map the feature vectors of each region Store in feature repository In the process, the region generation module is either a dense detection box generator for a one-stage object detection method or a region extraction network for a two-stage object detection method to generate candidate regions; the reliable feature vector selection in the sixth step is determined by the category prediction score of each interest region; if the highest category prediction score of the currently selected interest region bounding box is greater than the threshold θ, then the feature vector corresponding to the candidate box is considered a reliable feature vector. Step 7: Construct the contrastive learning loss function; where for unlabeled data... All prediction results If the predicted category is calibration frame features If a sample is taken, its corresponding contrastive learning loss function is in the form of: , , The number of sampling calibration boxes for the current image, where and These are sets of positive and negative sample feature vectors, respectively. The feature vectors in these sets come from a feature repository. The positive sample set consists of and Belonging to the same category The vector composition, It consists of vectors from all other categories; This is a vector similarity measurement method; the role of the contrastive learning loss function is to bring feature samples of the same category closer together in the high-dimensional feature space, while pulling samples of different categories further apart. Step 8, with labeled data Use the usual classification loss. and regression loss The total loss function for labeled data is used as the loss function during training. In unlabeled data The above adopts consistency loss. To optimize the model, the loss function is denoted as... The consistency method uses different transformations of an image as input and calculates the difference in predictions under different transformations; the total loss function for unlabeled data is... ; Step 9: Optimize the loss function using stochastic gradient descent, train the model, and repeat steps 5-8, iterating multiple times until the model's prediction results stabilize.

2. The method according to claim 1, wherein the specific functions and steps of each module in the second step include: I. Feature Extraction Module Image feature extraction: This module first preprocesses the training images, and then uses common feature extraction networks to extract features from the training images, including VGG, ResNet, MobileNet, RetinaNet, or EfficientNet; furthermore, the feature extraction network is combined with a feature processing network to further extract and optimize the image features, including BAM, CBAM, SPP, FPN, and / or NAS-FPN; II. Region Generation Module Candidate regions are generated, and these candidate regions, along with image features, serve as input to the output prediction module. III. Output Prediction Module Predict the type and location of the target; this module first performs RoI Pooling / Align on the features within the candidate region, which is the YOLO network for one-stage dense prediction, and the Faster RCNN network or the CenterNet network without anchor boxes for two-stage refined prediction.

3. The method according to claim 1, wherein in the fourth step, the feature mapping module The MLP (Multilayer Perceptron) is adopted, and its structure is: fully connected layer -> ReLU activation layer -> batch normalization -> fully connected layer.

4. The method according to claim 1, in the seventh step, for any two different feature vectors and Similarity measurement methods include: Using vector dot product MLP layer mapping Calculate the cosine distance between vectors or Gaussian method To measure the similarity between the two.

5. An information processing device, comprising a memory, a processor, and a program stored in the memory and executable on the processor, wherein executing the program implements the following characteristics: The computer program, when executed by a processor, implements the method as described in any one of claims 1 to 4.

Images