An extraterrestrial picture segmentation method and device based on semi-supervised learning

Abstract

The application discloses an extraterrestrial picture segmentation method and device based on semi-supervised learning, and the method comprises the following steps: acquiring a training data set and constructing a picture segmentation network; pre-training the picture segmentation network; feeding the sample picture into the pre-trained picture segmentation network; training the pre-trained picture segmentation network based on the comparison learning of the prediction segmentation result output by the first classifier and the down-sampling feature map output by the second classifier at the feature pixel level, wherein, after each round of training, the high-confidence unlabeled region and the labeled region are fused to obtain the labeled region for the next round of training; and inputting the target extraterrestrial picture into the trained picture segmentation network to obtain the segmentation result of the target extraterrestrial picture. The application introduces more supervision signals and designs the pixel-level comparison learning strategy, thereby significantly improving the picture segmentation performance.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of picture segmentation and the field of semi-supervised learning, and particularly relates to a method and device for extraterrestrial picture segmentation based on semi-supervised learning. BACKGROUND

[0002] Image segmentation is a technique and process of dividing an image into several specific regions with unique properties and presenting the target of interest. It is a key step from image processing to image analysis.

[0003] The existing image segmentation methods can be divided into the following four categories. A threshold-based segmentation method, which defines the output image based on the pixel value of the input image and the threshold defined by a person, is generally used for image segmentation with simple structure content. An edge-based segmentation method, which detects the edge by edge detection, detects the place where the gray level or structure has a mutation, and extracts the discontinuous edge region to determine the segmentation result. A clustering segmentation method based on histogram statistics, which determines the color and intensity of the pixels in the image, and counts the histogram result, is used to locate different clusters in the image as the segmentation result. A deep learning method based on front propagation, which designs a module with image segmentation function, is trained on a large data set, and only one parameter front propagation is needed to obtain the segmentation result in inference.

[0004] However, the first three methods cannot take advantage of large-scale labeled data sets, and the artificial design cannot simulate the complex calculation of the segmentation process, so the accuracy cannot meet the requirements in complex scene segmentation. The existing deep learning method based on front propagation cannot fully utilize the sparse labeling characteristics of extraterrestrial data, and the segmentation result still has a large space for improvement, which cannot meet the needs of practical applications. SUMMARY

[0005] To solve the above technical problems, the application provides a method and device for extraterrestrial picture segmentation based on semi-supervised learning, which can output the corresponding semantic segmentation map for a given extraterrestrial surface picture, has better performance in segmentation indicators, and can still produce reasonable segmentation results in the unlabeled region of the data set.

[0006] The technical scheme adopted by the application is as follows:

[0007] A method for extraterrestrial picture segmentation based on semi-supervised learning, comprising the following steps:

[0008] Obtain a training data set and construct a picture segmentation network, the sample extraterrestrial picture in the training data set comprises a labeled region and an unlabeled region, and the picture segmentation network comprises a feature extraction network, a feature output head and a feature down-sampling output head connected to the feature extraction network, a first classifier connected to the feature output head, and a second classifier connected to the feature down-sampling output head.

[0009] pre-training the picture segmentation network;

[0010] feeding the sample picture into the pre-trained picture segmentation network;

[0011] training the pre-trained picture segmentation network based on contrastive learning between the predicted segmentation result output by the first classifier and the down-sampled feature map output by the second classifier at a feature pixel level, wherein, after each round of training, high-confidence unlabeled regions are fused with labeled regions to obtain labeled regions for the next round of training;

[0012] inputting a target extraterrestrial picture into the trained picture segmentation network to obtain a segmentation result of the target extraterrestrial picture.

[0013] Further, the feature extraction network comprises M convolutional layers, wherein each convolutional layer is followed by a linear rectifier function, the first convolutional layer is followed by a pooling layer, the step length of a part of the convolutional layers is x, the step length of another part of the convolutional layers is 2x, x is a natural number, and the output of the feature extraction network is the fusion result of the output of the i-th convolutional layer and the output of the N-th convolutional layer, i < M.

[0014] Further, the feature output head comprises a multi-size aggregation pooling layer and a projection head.

[0015] Further, the projection head comprises a convolutional layer followed by batch normalization and a linear rectifier function.

[0016] Further, the pre-training of the picture segmentation network comprises:

[0017] based on the training data set, supervised learning is performed on the feature extraction network, the feature down-sampling output head, and the second classifier to obtain first parameters of the feature extraction network, the feature down-sampling output head, and the second classifier;

[0018] based on contrastive learning between the predicted segmentation result output by the first classifier and the down-sampled feature map output by the second classifier at a feature pixel level, the feature output head, the first classifier, and the feature extraction network, the feature down-sampling output head, and the second classifier with the first parameters are trained to obtain initial parameters of the picture segmentation network.

[0019] Further, the training of the feature output head, the first classifier, and the feature extraction network, the feature down-sampling output head, and the second classifier with the first parameters further comprises:

[0020] Based on the feature mean value of each category label region, a queue is maintained for each category, wherein the data in the queue is updated in a first-in first-out manner through the pre-trained picture segmentation network during the training process.

[0021] All the queues of all categories are collectively composed into a memory bank.

[0022] Further, the loss function for supervising the feature extraction network, the feature output head and the first classifier Wherein f i is the feature of the i-th position in the sample extraneous picture x output by the feature output head, c i is the true label of the i-th position of the sample, is the mean value of all features belonging to the c j category maintained in the memory bank, and tau is a temperature coefficient, and K is a kernel function.

[0023] Further, the loss function for training the pre-trained picture segmentation network is L = lambda s L s + L ce , wherein lambda s is a weight term, represents the probability of the i-th position in the feature f belonging to the category c j .

[0024] A storage medium, wherein the storage medium stores a computer program, and the computer program is configured to execute any of the above-mentioned methods when running.

[0025] An electronic device, comprising a memory and a processor, wherein the memory stores a computer program, and the processor is configured to execute the computer program to execute any of the above-mentioned methods.

[0026] The present application introduces more supervised signals through semi-supervised learning, and the designed pixel-level contrast learning strategy makes the feature space learned by the network more separable and robust, compared with the prior art, the present application significantly improves the picture segmentation performance, on the AI4MARS large-scale extraneous picture segmentation benchmark test set, the Frequency Weighted Intersection over Union index is improved from 72.67% to 88.82%, and the Mean Intersection over Union index is improved from 55.34% to 70.34%. BRIEF DESCRIPTION OF DRAWINGS

[0027] Figure 1A structural diagram of a picture segmentation neural network framework used by an embodiment of the present application.

[0028] Figure 2A An input picture of an embodiment of the present application.

[0029] Figure 2B A model prediction segmentation result picture of an embodiment of the present application. DETAILED DESCRIPTION

[0030] To make the above features and advantages of the present application more obvious and easy to understand, the following embodiments are specifically given, and are described in detail below with the aid of the accompanying drawings. It should be noted that the specific number of layers, the number of modules, the number of functions, and the settings of some layers given in the following embodiments are only a preferred implementation manner, and are not used for limitation, and those skilled in the art can select the number and set some layers according to actual needs, which should be understood.

[0031] The extraterrestrial picture segmentation method of the present application uses a training data set to train a picture segmentation neural network model; the neural network model includes a feature extraction network for extracting a feature representation of an input picture; a segmentation result output head taking the output of the feature extraction network as input; and a feature down-sampling module in parallel therewith, which down-samples the output features of the feature extraction network for contrast learning at the feature pixel level and introduces supervision information of unannotated parts.

[0032] The extraterrestrial picture segmentation method of the present application, as shown in Figure 1 , includes

[0033] Step 1: Collect a large number of extraterrestrial pictures and their corresponding segmentation labels to form a training data set, and due to the sparsity of annotation, there are a large number of unannotated pixel regions in the pictures.

[0034] Step 2: Build a neural network model for picture segmentation.

[0035] The network structure is as shown in Figure 1 , and the model is divided into a feature extraction network B seg , an output head H seg , and a feature down-sampling head , which are composed of three sub-networks.

[0036] The feature extraction network as a whole adopts a feature encoding part similar to ResNet-101, including 34 consecutive convolutional layers, each followed by a linear rectifier function (ReLU), and the first convolutional layer is followed by a 3x 3 max-pooling with a step size of 2 for down-sampling, and finally outputs a feature map down-sampled by 16 times. Among them, the step size of the 2nd-4th convolutional layer and the 32nd-34th convolutional layer of the original ResNet model is changed to 1, and the output features are the combination of the output content of the 4th convolutional layer and the last convolutional layer.

[0037] Output head H seg and feature down-sampling head Both of them use the Head part of DeepLabV3+, including a projection head, which is composed of a convolution layer followed by a batch normalization (Batch Normalization) and a linear rectifier function; a multi-size aggregation pooling layer, the input is processed by 4 dilated convolution layers with different dilation rates and an average pooling layer + convolution layer, then the features output by them are merged and input into a projection head to get the final result, which includes a convolution layer followed by a batch normalization and a linear rectifier function, and random deactivation (Dropout) is applied for regularization; finally, a classifier composed of convolution layers outputs the final prediction results or feature maps. seg The final output is the prediction probability of the network for the input belonging to different categories, and the channel number of the feature dimension is the total number of categories in the dataset; The final output of the network is the feature representation of the input sample, and the channel number of the feature dimension is 128.

[0038] The queue length of the memory bank maintained for contrastive learning is 32. To solve the problem of large computational overhead of pixel-level contrastive learning, a queue is maintained for each category to record the feature representation belonging to the category, which is given by the mean of the last N features belonging to the category label region. The feature representation of each category is maintained in a queue and updated in a first-in, first-out manner during the training process. All category queues together form a memory bank (Memory bank)

[0039] Step 3: Train the image segmentation neural network model.

[0040] Step 3.1 First pre-training model phase.

[0041] This phase is based on the training dataset, and the feature extraction network, feature output head and first classifier are supervised learning. Specifically, it includes:

[0042] 1. Feed the image data into the feature extraction network (baseline model) to combine the deep and shallow features of the network to obtain the feature representation of the image;

[0043] 2. Feed the features into the feature output head to get the segmentation results output by the network, which includes the predicted results of the labeled and unlabeled regions;

[0044] 3. The total loss function term of the first pre-training model phase is: L = L s , the inter-class contrastive learning loss function term

[0045] wherein, is the mean of all features belonging to class c j maintained in the memory bank, c i is the true label of the i-th position of sample x, τ is the temperature coefficient, usually 0.07, and K is the kernel function.

[0046] Step 3.2 Second pre-training phase model stage.

[0047] In this phase, the neural network model is trained based on the training data set. The total loss function term of this phase is: L = λ s L s + L ce , wherein λ s is the weight term, and λ s is usually set to 0.001.

[0048] L s is the same as the previous stage.

[0049]

[0050] wherein f is the feature map of sample x output by the neural network model, represents the probability of belonging to class c j in the i-th position of feature f in the predicted probability, c i is the true label of the i-th position of the sample.

[0051] Step 3.3 Third phase of model training.

[0052] The total loss function term of this phase is consistent with the second phase. However, unlike the second phase, the third phase adds the use of pseudo labels for unlabeled regions, pixel-level contrastive loss L s is optimized in the labeled region and the unlabeled region of the picture data. The pseudo label of the unlabeled region is generated in an online manner, and high-confidence regions selected according to a confidence threshold are used.

[0053] In an example, the above three stages are trained for 60 epochs.

[0054] Step 4: Inference phase,

[0055] In a test example, a reference picture to be tested is input (see Figure 2A ), and the final output is the corresponding semantic segmentation result (see Figure 2B ).

[0056] The above examples are only used to illustrate the technical solutions of the present application but not to limit the present application, and the ordinary skilled in the art can modify or equivalently replace the technical solutions of the present application without departing from the spirit and scope of the present application, and the protection scope of the present application should be subject to the description of the claims.

Claims

1. An extraterrestrial image segmentation method based on semi-supervised learning, comprising the steps of: obtaining a training data set, wherein sample extraterrestrial images in the training data set comprise labeled regions and unlabeled regions, and constructing an image segmentation network, wherein the image segmentation network comprises a feature extraction network, a feature output head and a feature down-sampling output head connected to the feature extraction network, a first classifier connected to the feature output head, and a second classifier connected to the feature down-sampling output head; pre-training the image segmentation network; feeding the sample extraterrestrial images into the pre-trained image segmentation network; training the pre-trained image segmentation network based on contrastive learning of predicted segmentation results output by the first classifier and down-sampled feature maps output by the second classifier at a feature pixel level, wherein high-confidence unlabeled regions are fused with labeled regions to obtain labeled regions for next round of training; and inputting a target extraterrestrial image into the trained image segmentation network to obtain a segmentation result of the target extraterrestrial image. The feature extraction network comprises M convolutional layers, wherein each convolutional layer is followed by a linear rectifier function, the first convolutional layer is followed by a pooling layer, the step length of a part of the convolutional layers is x, the step length of another part of the convolutional layers is 2x, x is a natural number, and the output of the feature extraction network is the fusion result of the output of the i-th convolutional layer and the output of the N-th convolutional layer, i < M. The feature output head comprises a multi-size aggregation pooling layer and a projection head. The projection head comprises a convolutional layer followed by batch normalization and a linear rectifier function. The pre-training of the image segmentation network comprises: based on the training data set, performing supervised learning on the feature extraction network, the feature down-sampling output head, and the second classifier to obtain first parameters of the feature extraction network, the feature down-sampling output head, and the second classifier; 2. The method of claim 1, wherein, based on contrastive learning of predicted segmentation results output by the first classifier and down-sampled feature maps output by the second classifier at a feature pixel level, training the feature output head, the first classifier, and the feature extraction network, the feature down-sampling output head, and the second classifier with the first parameters to obtain initial parameters of the image segmentation network.

3. The method of claim 1, wherein, The training of the feature output head, the first classifier, and the feature extraction network, the feature down-sampling output head, and the second classifier with the first parameters further comprises:

4. The method of claim 3, wherein, based on the feature mean of each category label region, maintaining a queue for each category, wherein the data in the queue is updated in a first-in-first-out manner during the training process of the pre-trained image segmentation network.

5. The method of claim 1, wherein, All queues of all categories are combined to form a memory bank. The computer program is configured to execute the method of any one of claims 1-8 when running. ​ 6. The method of claim 5, wherein, ​ ​ ​ 7. The method of claim 6, wherein, The loss function for supervising the feature extraction network, the feature output head, and the first classifier wherein f i is the feature of the i-th position in the sample out-of-distribution picture x output by the feature output head, c i is the true label of the i-th position of the sample, is the mean of all features belonging to the c j category maintained in the memory bank, is the mean of all features belonging to the c i category maintained in the memory bank, and τ is a temperature coefficient, and K is a kernel function.

8. The method of claim 7, wherein, The loss function L for training the pre-trained picture segmentation network is s L s +L ce , wherein λ s is a weight term, represents the probability of belonging to class c j in the prediction probability of the i-th position in feature f, represents the probability of belonging to class c j in the prediction probability of the i-th position in feature f.

9. A storage medium having stored therein a computer program, wherein, ​ 10. An electronic device comprising a memory having a computer program stored therein and a processor arranged to run the computer program to perform the method of any one of claims 1-8.

Images