Training method and device of image segmentation model

Abstract

The invention discloses a training method and device of an image segmentation model, and relates to the technical field of computers. A specific embodiment of the method comprises the steps of obtaining a first image sample and a bounding box of a target object in the first image sample, and generating a first training sample; training the bounding box detection capability of the image segmentation model by using the first training sample to obtain a first image segmentation model; performing image bounding box labeling on the second image sample by using the first image segmentation model, and training the image segmentation capability of the first image segmentation model by using a labeling result and a general image segmentation model to obtain a second image segmentation model; and finally, carrying out overall optimization training on the second image segmentation model, wherein the scale of the image segmentation model is smaller than the specification of a general image segmentation model. According to the embodiment of the invention, the training of the lightweight image segmentation model is realized, the marking time and cost of the second image sample are saved, and the segmentation performance of the small-scale image segmentation model is ensured.

Description

Technical Field

[0001] The present invention relates to the field of computer technology, and in particular to a training method and device for an image segmentation model. Background Art

[0002] Image segmentation is a key technology in the field of computer vision. It aims to identify and segment each independent object instance from an image. It not only needs to distinguish between the foreground and background, but also needs to distinguish between different object instances. Existing image segmentation technologies are mainly based on image instance segmentation models built using deep learning methods.

[0003] In the process of implementing the present invention, the inventors found that the prior art has the following problems:

[0004] Existing image segmentation methods require a lot of computing resources and time during the training process. The demand for computing resources is extremely high, making it difficult for developers with limited resources to fully train and verify the model, which is not conducive to the practical application and implementation of the model. Summary of the invention

[0005] In view of this, an embodiment of the present invention provides a method and device for training an image segmentation model, which is based on a small-scale image segmentation model and utilizes the bounding box detection capability of the image segmentation model and a general image segmentation model to realize automatic labeling of the second image sample, thereby saving labeling time and cost, and further lightweight training the image segmentation capability of the image segmentation model, and then through overall optimization training of the model's bounding box detection capability and image segmentation capability, the segmentation performance of the small-scale image segmentation model is guaranteed, thereby realizing that the trained image segmentation model has good segmentation performance without occupying too many resources.

[0006] To achieve the above object, according to one aspect of an embodiment of the present invention, a method for training an image segmentation model is provided, comprising:

[0007] Acquire a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; use the first training sample to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model;

[0008] Using the first image segmentation model to annotate the image bounding box of the second image sample, and then using the annotated result and the general image segmentation model to train the image segmentation capability of the first image segmentation model to obtain a second image segmentation model;

[0009] Based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than that of the general image segmentation model.

[0010] Optionally, the second image sample is annotated with an image bounding box using the first image segmentation model, and the image segmentation capability of the first image segmentation model is trained using the annotation result and the general image segmentation model, including: inputting the second image sample into the first image segmentation model to obtain the bounding box of the target object in the second image sample and the corresponding segmented image output by the first image segmentation model; inputting the second image sample and the bounding box of the target object in the second image sample into the general image segmentation model to obtain a standard segmented image output by the general image segmentation model; and training the image segmentation capability of the first image segmentation model by calculating the minimum loss function between the standard segmented image and the segmented image.

[0011] Optionally, based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained, including: inputting a third image sample into the second image segmentation model, and utilizing the bounding box detection capability and image segmentation capability of the second image segmentation model to perform bounding box annotation and segmented image annotation on target objects in the third image sample; screening the bounding boxes and corresponding segmented images of the annotated target objects in the third image sample, and generating a second training sample based on the screening result; and using the second training sample to train the second image segmentation model.

[0012] Optionally, before obtaining the first image sample and the bounding box of the target object annotated in the first image sample, the method further includes: annotating the bounding box of the target object in the first image sample according to the annotation rules of a general image segmentation model, wherein the annotation rule is to use a specified mark at a specified position to characterize the position of the target object in the first image sample.

[0013] According to a second aspect of an embodiment of the present invention, there is provided a method for image segmentation, comprising:

[0014] Obtain the image to be segmented;

[0015] The image to be segmented is input into a pre-constructed image segmentation model, and a bounding box and a segmented image of the target object output by the image segmentation model are obtained; wherein the bounding box indicates the position of the target object in the image to be segmented, the segmented image is segmented from the image to be segmented and contains the target object, and the image segmentation model is constructed according to any one of the training methods provided in the first aspect of an embodiment of the present invention.

[0016] Optionally, the image to be segmented is input into a pre-built image segmentation model to obtain a bounding box and a segmented image of a target object output by the image segmentation model, including: inputting the image to be segmented into the image segmentation model, performing feature extraction on the image to be segmented through a backbone network of the image segmentation model to obtain image features of the image to be segmented; performing bounding box detection of the target object on the image features through a detection network of the image segmentation model to obtain a bounding box of the target object output by the image segmentation model; performing image segmentation of the target object on the image features through a segmentation network of the image segmentation model to obtain a segmented image of the target object output by the image segmentation model.

[0017] According to a third aspect of an embodiment of the present invention, there is provided an image annotation method, comprising:

[0018] Get the image to be annotated;

[0019] The image to be annotated is input into a pre-constructed image segmentation model, and the image to be annotated is annotated according to the bounding box of the target object and the corresponding segmented image output by the image segmentation model, wherein the image segmentation model is constructed according to any training method provided in the first aspect of an embodiment of the present invention.

[0020] According to a fourth aspect of an embodiment of the present invention, a training device for an image segmentation model is provided, comprising:

[0021] A first training module is used to obtain a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; and use the first training sample to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model;

[0022] A second training module is used to use the first image segmentation model to annotate the image bounding box of the second image sample, and then use the annotation result and the general image segmentation model to train the image segmentation ability of the first image segmentation model to obtain a second image segmentation model;

[0023] The third training module is used to perform overall optimization training on the second image segmentation model based on the bounding box detection capability and image segmentation capability of the second image segmentation model to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than that of the general image segmentation model.

[0024] According to a fifth aspect of an embodiment of the present invention, there is provided an electronic training device for an image segmentation model, comprising:

[0025] one or more processors;

[0026] a storage device for storing one or more programs,

[0027] When the one or more programs are executed by the one or more processors, the one or more processors implement the method provided in the first aspect of the embodiment of the present invention.

[0028] According to a sixth aspect of an embodiment of the present invention, there is provided a computer-readable medium having a computer program stored thereon, wherein the program, when executed by a processor, implements the method provided in the first aspect of the embodiment of the present invention.

[0029] According to a seventh aspect of an embodiment of the present invention, there is provided a computer program product, comprising a computer program, wherein when the computer program is executed by a processor, the method provided in the first aspect of the embodiment of the present invention is implemented.

[0030] An embodiment of the invention has the following advantages or beneficial effects: a first training sample is generated by acquiring a first image sample and a bounding box of a target object annotated in the first image sample; the first training sample is used to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model; the first image segmentation model is used to annotate the image bounding box of a second image sample, and the annotated result and a general image segmentation model are used to train the image segmentation capability of the first image segmentation model to obtain a second image segmentation model; based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained to obtain a trained image segmentation model, and the scale of the image segmentation model is smaller than that of the general image segmentation model. A technical solution is provided, which realizes a lightweight image segmentation model training method, and the bounding box detection capability of the image segmentation model and the general image segmentation model are used to realize automatic annotation of the second image sample, saving annotation time and cost, and then the segmentation performance of the small-scale image segmentation model is guaranteed by overall optimization training of the bounding box detection capability and image segmentation capability of the model, thereby realizing that the trained image segmentation model has good segmentation performance without occupying too many resources. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] The accompanying drawings are used to better understand the present invention and do not constitute an improper limitation of the present invention.

[0032] Figure 1 is a schematic diagram of the main process of the training method of the image segmentation model according to an embodiment of the present invention;

[0033] Figure 2 is a detailed flow chart of a training method for an image segmentation model according to a reference embodiment of the present invention;

[0034] Figure 32 is a schematic diagram showing the principle of obtaining a second image segmentation model by training with a general image segmentation model according to an embodiment of the present invention;

[0035] Figure 4 is a schematic diagram of main modules of a training device for an image segmentation model according to an embodiment of the present invention;

[0036] Figure 5 is an exemplary system architecture diagram to which embodiments of the present invention may be applied;

[0037] Figure 6 It is a schematic diagram of the structure of a computer system of a terminal device or a server suitable for implementing an embodiment of the present invention. DETAILED DESCRIPTION

[0038] It should be noted that the acquisition, storage and application of user personal information involved in the technical solution of the present disclosure are in compliance with the provisions of relevant laws and regulations and do not violate public order and good morals.

[0039] The following is a description of exemplary embodiments of the present invention in conjunction with the accompanying drawings, including various details of the embodiments of the present invention to facilitate understanding, which should be considered as merely exemplary. Therefore, it should be recognized by those of ordinary skill in the art that various changes and modifications may be made to the embodiments described herein without departing from the scope and spirit of the present invention. Similarly, for clarity and conciseness, the description of well-known functions and structures is omitted in the following description.

[0040] Existing image segmentation methods require a lot of computing resources and time during the training process. The demand for computing resources is extremely high, making it difficult for developers with limited resources to fully train and verify the model, which is not conducive to the practical application and implementation of the model.

[0041] In order to solve the above-mentioned problems existing in the prior art, the present invention proposes a training method for an image segmentation model, uses a first training sample to train the bounding box detection capability of the image segmentation model, and uses the bounding box detection capability of the image segmentation model and a general image segmentation model to train the image segmentation capability of the image segmentation model, and then obtains the trained image segmentation model through overall optimization training of the bounding box detection capability and the image segmentation capability, thereby realizing a lightweight training method for the image segmentation model, and uses the bounding box detection capability of the image segmentation model and the general image segmentation model to realize automatic labeling of the second image sample, thereby saving labeling time and cost, and then ensures the segmentation performance of the small-scale image segmentation model through overall optimization training of the bounding box detection capability and the image segmentation capability of the model, thereby realizing that the trained image segmentation model has good segmentation performance without occupying too many resources.

[0042] Figure 1is a schematic diagram of the main process of the training method of the image segmentation model according to an embodiment of the present invention, such as Figure 1 As shown, the image segmentation method of the embodiment of the present invention includes the following steps S101 to S103.

[0043] Step S101, obtaining a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; using the first training sample to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model.

[0044] Specifically, the image segmentation model is used to identify each object instance from the image. The image segmentation model can be applied to a variety of business fields, such as the image segmentation of the driving environment image in the intelligent driving business field, segmenting the pedestrian images therein; the image segmentation of the cargo stacking images on the turnover vehicle in the logistics loading and unloading business field, segmenting the individual cargo images, etc.

[0045] When the model training system receives a training task for an image segmentation model, it parses the training task to obtain sample information such as the specific business scenario in which the image segmentation model is applied and sample requirements (such as the number of samples, sample sources, etc.). Based on the obtained sample information, the first image sample that meets the requirements is extracted from the corresponding data file. The specific extraction method is to extract samples according to the extraction requirements if there are clear requirements in the sample requirements; if there are no clear requirements, random extraction can be selected or the corresponding sample extraction tasks can be allocated according to factors such as different categories and different times. The embodiment of the present invention does not make specific limitations.

[0046] Furthermore, based on the acquired first image samples, the bounding boxes of the target objects in each image in the first image samples are annotated. Since the purpose of this step is to construct the first image samples with bounding boxes as the first training samples, it is only necessary to annotate the bounding boxes of the first image samples. The specific annotation can be directly based on the contour of the target object, and the first image samples annotated with bounding boxes constitute the first training samples.

[0047] Each training image in the first training sample is input into the image segmentation model preset in the embodiment of the present invention, and the bounding box of the target object in the training image output by the image segmentation model is obtained. The bounding box of the target object annotated in the training image is then used to train the bounding box detection capability of the image segmentation model, and the minimum loss function between the bounding box output by the model and the annotated bounding box of the corresponding training image is calculated. When the training converges, a first image segmentation model with bounding box detection capability is obtained.

[0048] Step S102: using the first image segmentation model to annotate the image bounding box of the second image sample, and then using the annotation result and the general image segmentation model to train the image segmentation capability of the first image segmentation model to obtain the second image segmentation model.

[0049] Specifically, from the above-mentioned embodiments, it can be known that the first image segmentation model is essentially an image segmentation model with bounding box detection capability, and the image segmentation model also needs to have image segmentation capability, that is, it can segment the outline of the target object from the received image, so it is necessary to train the image segmentation capability of the first image segmentation model. It can be understood that the bounding box annotation of the target object in the image sample only needs to annotate its outline, while the segmented image annotation of the target object is more complicated than the bounding box annotation. Of course, it can also be implemented using some extraction tools, but when the number of image samples is usually large, it still takes a certain amount of time and effort to annotate the segmented image.

[0050] In order to solve the above problems and reduce the workload of annotation, the embodiment of the present invention uses a general image segmentation model SAM (Segment Anything Model). SAM is a multimodal model that performs image segmentation through prompt information, and its prompt information may include bounding box annotation, mouse click, etc. Although the general image segmentation model SAM can provide a solution for image segmentation, the model is a general model that relies on a large image encoder ViT, and it also requires certain computing resources to support it, which is not conducive to saving resource overhead and cannot solve the problem of excessive resource occupation in the prior art. However, the embodiment of the present invention can consider using the image segmentation capability of the general image segmentation model SAM only in the training stage, using the trained SAM model with fixed parameters to alleviate the training resource consumption of the image segmentation model, and using the SAM model to annotate the segmented image for the second image sample. That is, the bounding box of the target object annotated by the first image segmentation model in the second image sample is used as the prompt information of the SAM model, so as to obtain the segmented image of the target object in the second image sample, and then train the image segmentation capability of the first image segmentation model according to the segmented image output by the general image segmentation model SAM, and obtain the second image segmentation model when the training converges.

[0051] It can be understood that the second image segmentation model of the embodiment of the present invention is essentially an image segmentation model with bounding box detection capabilities and image segmentation capabilities.

[0052] Step S103: Based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than that of the general image segmentation model.

[0053] Specifically, after the above-mentioned training of the image segmentation model's bounding box detection capability and image segmentation capability respectively, a second image segmentation model is obtained. Considering that the second image segmentation model is obtained by performing bounding box detection and image segmentation training respectively, and in actual applications, bounding box detection and image segmentation are parallel, in order to ensure the performance of the model, it is necessary to further perform overall optimization training on the bounding box detection capability and image segmentation capability of the second image segmentation model, and use the trained image segmentation model as the result of this training task.

[0054] It should be noted that the image segmentation model constructed in the embodiment of the present invention is not highly dependent on the large image encoder ViT because it has been trained with the first training sample and the general image segmentation model, and a conventional image encoder can be used. Accordingly, the scale of the image segmentation model is smaller than that of the general graphic segmentation model. After the image segmentation model is trained using the SAM model, the SAM model can be removed, thereby saving system resource consumption.

[0055] It is understandable that the performance of the image segmentation model essentially depends on the quality of the training samples, and the quality of the training samples depends on the quality of the annotations. A large number of rich and accurately annotated training samples can effectively guarantee the image segmentation training effect of the model. That is, the performance of the model is highly dependent on the training sample set. High-quality annotation data is time-consuming and labor-intensive, and also limits the rapid deployment of the model in new fields. The embodiment of the present invention uses the general image segmentation model SAM to annotate the segmented image of the target object of the second image sample, which can greatly reduce the sample annotation cost, improve the sample annotation efficiency, and thus speed up the deployment of the image segmentation model in new fields, which is more conducive to the promotion and application of the image segmentation model.

[0056] Figure 2 FIG. 1 is a detailed flow chart of a training method for an image segmentation model according to a reference embodiment of the present invention. As another embodiment of the present invention, Figure 2 As shown, the training method of the image segmentation model includes:

[0057] Step S201, obtaining a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; using the first training sample to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model.

[0058] According to one embodiment of the present invention, before obtaining a first image sample and a bounding box of a target object annotated in the first image sample, the method further includes: annotating the bounding box of the target object in the first image sample according to an annotation rule of a general image segmentation model, wherein the annotation rule is to use a specified mark at a specified position to characterize the position of the target object in the first image sample.

[0059] Specifically, when annotating the bounding box of the target object in the first image sample, in view of the fact that the bounding box annotated by the embodiment of the present invention is mainly used to provide prompt information for the general image segmentation model SAM, it is preferred to annotate the bounding box of the target object according to the annotation rules of the general image segmentation model SAM. Compared with the traditional direct annotation of the bounding box of the target object, the annotation rules of the general image segmentation model SAM can greatly reduce the annotation cost, because the prompt information of the general image segmentation model SAM not only supports box selection, but also supports point selection, that is, only the mark generated by clicking the mouse at the location of the target object can be completed. Annotation of the target object. For example, for a rectangular packaging box of the target object, the conventional annotation method needs to annotate 4 points of the packaging box, but using the annotation rules of the SAM model only requires annotating one point within the area of ​​the target object.

[0060] Step S202: input the second image sample into the first image segmentation model, and obtain the bounding box of the target object in the second image sample and the corresponding segmented image output by the first image segmentation model.

[0061] Specifically, the first training sample is mainly used to train the bounding box detection capability of the image segmentation model. Before the second image sample is input into the first image segmentation model, the second image sample for training the image segmentation capability of the model is screened out. The second image sample is then input into the first image segmentation model. Since the first image segmentation model comes from the image segmentation model, and the image segmentation model includes two modules, namely, bounding box detection and image segmentation, the first image segmentation model can output the bounding box and the corresponding segmented image of the target object in the second image sample. However, the segmented image output at this time may be significantly different from the actual segmented image of the target object.

[0062] The embodiment of the present invention uses the first image segmentation model to obtain the bounding box of the target object in the second image sample to realize automatic bounding box annotation of the second image sample. The first image segmentation model is used to complete the automatic bounding box annotation of the second image sample, thereby improving the annotation efficiency and saving the annotation cost.

[0063] Step S203: input the second image sample and the bounding box of the target object in the second image sample into a general image segmentation model to obtain a standard segmentation image output by the general image segmentation model.

[0064] Specifically, the embodiment of the present invention uses the bounding box of the target object annotated in the second image sample as the prompt information of the general image segmentation model SAM, and also inputs the image sample corresponding to the target object in the second image sample into the general image segmentation model SAM. The parameters of the SAM model used in the embodiment of the present invention are fixed, and the image segmentation ability of the SAM model is used to segment the standard segmentation image of the target object from the image sample. We use the standard segmentation image as the true value, which is essentially to use the SAM model to perform image segmentation and annotation on the target object in the second image sample, thereby realizing the automation of image segmentation and annotation.

[0065] Step S204: training the image segmentation capability of the first image segmentation model by calculating the minimum loss function between the standard segmented image and the segmented image, thereby obtaining a second image segmentation model.

[0066] Specifically, the standard segmentation image output by the general image segmentation model SAM is taken as the true value, and is regarded as the segmentation image of the target object marked in the second image sample. The image segmentation ability of the first image segmentation model is trained by calculating the minimum loss function between the segmentation image output by the first image segmentation model and the corresponding standard segmentation image output by SAM, and after the training converges, the second image segmentation model is obtained.

[0067] It can be understood that the embodiment of the present invention collectively refers to the training process of the first image segmentation model that provides bounding box annotations for the SAM prompt information and the training process of the second image segmentation model that is directly trained by the SAM model as the first stage training. The first stage training needs to be implemented in conjunction with the SAM model, which is equivalent to the pre-training stage of the image segmentation model.

[0068] The pre-training of the image segmentation model in the embodiment of the present invention makes use of the general image segmentation model SAM. It only needs to perform bounding box annotation on the first image sample to realize the training of the bounding box detection capability and image segmentation capability of the image segmentation model. That is, only a small amount of annotation is needed to drive the pre-training of the entire image segmentation model, which greatly improves the efficiency of pre-training, reduces resource overhead, and facilitates the promotion and application of the model.

[0069] Step S205: input the third image sample into the second image segmentation model, and use the bounding box detection capability and image segmentation capability of the second image segmentation model to perform bounding box annotation and segmented image annotation on the target object in the third image sample.

[0070] Specifically, in order to ensure the performance of the second image segmentation model, the embodiment of the present invention also needs to screen out the third image sample again. The role of the third image sample is to conduct comprehensive and overall training on the bounding box detection capability and image segmentation capability of the second image segmentation model. In addition, considering that the second image segmentation model at this time has undergone bounding box detection training and image segmentation training, the second image segmentation model has certain bounding box detection capabilities and image segmentation capabilities. In order to further save the construction consumption of the training sample set and improve the construction efficiency of the training sample, the embodiment of the present invention uses the second image segmentation model to automatically generate candidate samples for the second training sample for independent training in the next stage. First, the third image sample is self-annotated using the second image segmentation model, that is, the third image sample is input into the second image segmentation model, and the bounding box and segmented image of the target object in the third image sample output by the second image segmentation model are used to implement the bounding box annotation and segmented image annotation of the third image sample.

[0071] Step S206: Screen the bounding box of the target object marked in the third image sample and the corresponding segmented image, and generate a second training sample according to the screening result.

[0072] Specifically, since the second image segmentation model is only the result of pre-training, in order to ensure the quality of the training samples, the bounding box annotations and segmented image annotations of the third image sample using the second image segmentation model are screened. The annotated bounding boxes and segmented images and the corresponding image samples are submitted to the annotation review system. The annotation review system reviews and screens the annotation results according to the configured screening criteria, selects the annotation results with high accuracy, and uses these screened annotation results to construct the second training sample. It can be understood that if the number of screened annotation results does not meet the number requirements of the second training sample set, the annotation results can be data enhanced to meet the number requirements of the second training sample.

[0073] Step S207: Use the second training sample to train the second image segmentation model to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than the scale of the general graphic segmentation model.

[0074] Finally, the second image segmentation model is independently trained using the second training samples constructed above. The participation of the general image segmentation model SAM is no longer required, so that the second image segmentation model can be independently trained. When the training converges, the trained image segmentation model is obtained.

[0075] It can be understood that the embodiment of the present invention refers to the above-mentioned training process of the second image segmentation model as the second stage training. The second stage training does not need to be combined with the SAM model. The second image segmentation model only constructs candidate samples of the second training samples, and then selects samples that meet the requirements to generate second training samples, and uses the second training samples to independently perform overall optimization training on the model.

[0076] Based on the above-mentioned training method of the image segmentation model, an embodiment of the present invention further provides a method for image segmentation, which may specifically include: obtaining an image to be segmented; inputting the image to be segmented into a pre-constructed image segmentation model, and obtaining a bounding box and a segmented image of the target object output by the image segmentation model; wherein the bounding box indicates the position of the target object in the image to be segmented, the segmented image is segmented from the image to be segmented and contains the target object, and the image segmentation model is constructed according to any one of the above-mentioned training methods of the image segmentation model.

[0077] Specifically, the image segmentation method provided by the embodiment of the present invention can be applied to a variety of business fields, such as performing image segmentation on images in the driving environment in the intelligent driving business field to segment out pedestrian images therein, performing image segmentation on cargo stacking images on turnover vehicles in the logistics loading and unloading business field to segment out individual cargo images, etc. When receiving an image segmentation task, the image segmentation system obtains the image to be segmented therein by parsing the image segmentation task, or obtains the address where the image to be segmented is stored by parsing the image segmentation task, and obtains the image to be segmented by accessing the address. The content of the image to be segmented is as described above, which can be a road environment image, a cargo stacking image, or an internal structure image of a product, etc.

[0078] Furthermore, in order to ensure the image segmentation effect and performance of the image segmentation model, the embodiment of the present invention adopts a two-stage training method, which utilizes the general image segmentation model SAM, which has excellent image segmentation ability when the bounding box of the image to be segmented is known to be annotated. The SAM model is first combined and the SAM model is used as a teacher model. The image segmentation model learns the image segmentation ability of SAM and performs the first stage of joint training; then the constructed second training sample is used to independently train the image segmentation second model to improve the image segmentation accuracy of the trained image segmentation model.

[0079] According to one embodiment of the present invention, the image to be segmented is input into a pre-built image segmentation model, and a bounding box and a segmented image of a target object output by the image segmentation model are obtained, including: inputting the image to be segmented into the image segmentation model, performing feature extraction on the image to be segmented through a backbone network of the image segmentation model, and obtaining image features of the image to be segmented; performing bounding box detection of the target object on the image features through a detection network of the image segmentation model, and obtaining the bounding box of the target object output by the image segmentation model; performing image segmentation of the target object on the image features through a segmentation network of the image segmentation model, and obtaining a segmented image of the target object output by the image segmentation model.

[0080] Specifically, the image segmentation model of the embodiment of the present invention includes a backbone network, a detection network and a segmentation network. The number of layers of the backbone network, the detection network and the segmentation network are all limited within a specified layer number threshold and not less than the lower limit number of layers of each network, so as to take into account the image segmentation capability and model scale of the model. Furthermore, the image to be segmented is input into a pre-built image segmentation model, and the backbone network of the image segmentation model first extracts features of the image to be segmented to obtain image features of the image to be segmented; then the image features are shared with the detection network and the segmentation network within the model, and the detection network and the segmentation network respectively perform bounding box detection and image segmentation of the target object on the image to be segmented according to the received image features and the target object, and output the bounding box and segmented image of the target object.

[0081] Since the construction of the image segmentation model in the embodiment of the present invention takes the model scale factor into consideration, the scale of the established image segmentation model is a small-scale model under the premise of ensuring the image segmentation capability, thereby saving system resources, thereby balancing the image segmentation and real-time processing capabilities of the model, and realizing lightweight image segmentation.

[0082] It can be understood that the bounding box of the target object output by the image segmentation model is the outer shape box of the target object in the image to be segmented. The bounding box not only indicates the position of the target object in the image to be segmented, but also shows the size of the target object in the image to be segmented, so as to better meet the image segmentation requirements. The segmented image of the target object is the complete image representation of the target object in the image to be segmented, which is used to more vividly reflect the individual characteristics of the target object.

[0083] According to one embodiment of the present invention, after obtaining the bounding box and segmented image of the target object output by the image segmentation model, the method further includes: using the bounding box and the segmented image of the target object, and the image to be segmented, to update the second training sample of the image segmentation model to iteratively update the image segmentation model.

[0084] It can be understood that in order to improve the performance and generalization ability of the image segmentation model to better serve the needs of image segmentation, after obtaining the bounding box and segmented image of the target object output by the image segmentation model, the user's feedback on the bounding box and segmented image output by the model is used to determine whether to use the output of the image segmentation model for the second training sample that can be used for model iterative training. If the result of the feedback model output is not ideal, then ignore the model output result and try to add high-quality annotated data for this business scenario to train the image segmentation model again; if the score of the result of the user's feedback model output is higher than the preset score threshold, then the bounding box and segmented image of the target object, as well as the image to be segmented, are added to the second training sample, so as to iteratively update the image segmentation model and further improve the performance of the image segmentation model.

[0085] Based on the above-mentioned training method of the image segmentation model, an embodiment of the present invention further provides an image annotation method, which may specifically include: obtaining an image to be annotated; inputting the image to be annotated into a pre-constructed image segmentation model, and annotating the image to be annotated according to the bounding box of the target object and the corresponding segmented image output by the image segmentation model, wherein the image segmentation model is constructed according to any of the above-mentioned training methods of the image segmentation model.

[0086] Specifically, the image segmentation results of the image to be segmented by the image segmentation model constructed through the above training can be known, and the image segmentation results include the bounding box of the target object and the segmented image of the target object. Considering that in the construction process of the training sample of the model, the most time-consuming and laborious link is the labeling step of the individual image in the image, it is usually necessary to perform manual work to label the borders and segmented images of the individual objects required in the massive images. Then the image segmentation model of the embodiment of the present invention can also be used for image labeling tasks in addition to image segmentation tasks. Therefore, when an image labeling request is received, the image segmentation model is used for labeling, and the image to be labeled is input into the pre-constructed image segmentation model. It can be understood that at this time, it is also necessary to clarify the individual objects or target objects that need to be labeled to the image segmentation model, and then according to the bounding box of the target object output by the image segmentation model and the corresponding segmented image, the automatic labeling of the image to be labeled is realized, which greatly saves labor costs and the efficiency of labeling, and accordingly, the construction efficiency of the training sample set is also improved.

[0087] Figure 3It is a schematic diagram of the principle of the second image segmentation model obtained by training with a general image segmentation model in an embodiment of the present invention. The general image segmentation model SAM in an embodiment of the present invention mainly includes an image encoder ViT, image feature embedding, and mask decoding based on prompt information; the image segmentation model includes a backbone network, a detection network, and a segmentation network. The backbone network is mainly used for image feature extraction, the detection network is used to detect the bounding box of the target object based on the image features, and the segmentation network is used to segment the image of the target object based on the image features. It can be understood that the first image segmentation model is obtained by training the image segmentation model for bounding box detection capabilities, and the composition structure of the model will not change. The first image segmentation model also includes a backbone network, a detection network, and a segmentation network.

[0088] Specifically, first obtain a first image sample and a bounding box of a target object annotated in the first image sample, and generate a first training sample for training the bounding box detection capability of an image segmentation model. Use the first training sample to train the bounding box detection capability of the image segmentation model, and define the trained model as the first image segmentation model. Then obtain a second image sample, input the second image sample into the first image segmentation model and the general image segmentation model SAM, use the first image segmentation model to annotate the bounding box Bbox of the target object in the second image sample, and input Bbox as prompt information into the mask decoding module based on the prompt information of SAM, and obtain the standard segmentation image TMask output by the SAM model. Finally, the first image segmentation model is trained by calculating the minimum loss function Loss between the standard segmentation image TMask and the segmentation image Mask output by the first image segmentation model, and when the training converges, the second image segmentation model is obtained.

[0089] The embodiment of the present invention uses the general image segmentation model SAM to train a custom image segmentation model with a simple architecture and a small scale, which not only alleviates the training intensity of the model, but also can drive the training of the entire image segmentation model based on a small amount of sample annotations (first training samples), greatly improving the efficiency of constructing training samples, and also ensuring the quality of training samples, thereby greatly improving the flexibility of image segmentation model application and the efficiency of implementation. In addition, the image segmentation model trained in the embodiment of the present invention can not only be used to segment images, but also can be used to automatically annotate images, expanding the application field and adaptability of the model.

[0090] Figure 4 Schematic diagram of the main modules of the training device of the image segmentation model according to the embodiment of the present invention. Figure 4 As shown, the training device 400 of the image segmentation model mainly includes a first training module 401, a second training module 402 and a third training module 403.

[0091] The first training module 401 is used to obtain a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; and use the first training sample to train the bounding box detection capability of the image segmentation model to obtain a first image segmentation model;

[0092] A second training module 402 is used to use the first image segmentation model to annotate the image bounding box of the second image sample, and then use the annotation result and the general image segmentation model to train the image segmentation ability of the first image segmentation model to obtain a second image segmentation model;

[0093] The third training module 403 is used to perform overall optimization training on the second image segmentation model based on the bounding box detection capability and image segmentation capability of the second image segmentation model to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than that of the general image segmentation model.

[0094] According to one embodiment of the present invention, the second training module 402 is further used to: input the second image sample into the first image segmentation model to obtain the bounding box of the target object in the second image sample and the corresponding segmented image output by the first image segmentation model; input the second image sample and the bounding box of the target object in the second image sample into a general image segmentation model to obtain a standard segmented image output by the general image segmentation model; and train the image segmentation capability of the first image segmentation model by calculating the minimum loss function between the standard segmented image and the segmented image.

[0095] According to another embodiment of the present invention, the third training module 403 is further used to: input the third image sample into the second image segmentation model, and use the bounding box detection capability and image segmentation capability of the second image segmentation model to perform bounding box annotation and segmented image annotation on the target object in the third image sample; screen the bounding box and the corresponding segmented image of the marked target object in the third image sample, and generate a second training sample based on the screening result; and use the second training sample to train the second image segmentation model.

[0096] According to another embodiment of the present invention, the training device 400 of the image segmentation model also includes a labeling module (not shown in the figure), which is used to: before obtaining the first image sample and the bounding box of the target object annotated in the first image sample, label the bounding box of the target object in the first image sample according to the labeling rules of the general image segmentation model, wherein the labeling rule is to use a specified mark at a specified position to represent the position of the target object in the first image sample.

[0097] Figure 5is an exemplary system architecture diagram to which embodiments of the present invention can be applied.

[0098] like Figure 5 As shown, system architecture 500 may include terminal devices 501, 502, 503, a network 504 and a server 505. Network 504 is used to provide a medium for communication links between terminal devices 501, 502, 503 and server 505. Network 504 may include various connection types, such as wired, wireless communication links or optical fiber cables, etc.

[0099] Users can use terminal devices 501, 502, 503 to interact with server 505 through network 504 to receive or send messages, etc. Various communication client applications, such as image segmentation applications, etc., can be installed on terminal devices 501, 502, 503 (only as an example).

[0100] The terminal devices 501 , 502 , and 503 may be various electronic devices having a display screen and supporting web browsing, including but not limited to smart phones, tablet computers, laptop computers, desktop computers, and the like.

[0101] The server 505 may be a server that provides various services, such as a background management server that provides support for the training of the image segmentation model performed by the user using the terminal devices 501, 502, and 503 (for example only). The background management server may obtain a first image sample and a bounding box of a target object annotated in the first image sample to generate a first training sample; use the first training sample to train the bounding box detection capability of the image segmentation model to obtain a first image segmentation model; use the first image segmentation model to annotate the image bounding box of the second image sample, and then use the annotation result and the general image segmentation model to train the image segmentation capability of the first image segmentation model to obtain a second image segmentation model; based on the bounding box detection capability and image segmentation capability of the second image segmentation model, perform overall optimization training on the second image segmentation model to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than the scale of the general image segmentation model, and the processing result (for example, the trained image segmentation model - for example only) is fed back to the terminal device.

[0102] It should be noted that the training method of the image segmentation model provided in the embodiment of the present invention is generally executed by the server 505 , and accordingly, the training device of the image segmentation model is generally set in the server 505 .

[0103] It should be understood that Figure 5 The number of terminal devices, networks and servers in the embodiment is only for illustration. Any number of terminal devices, networks and servers may be provided according to implementation requirements.

[0104] Reference below Figure 6 , Figure 6 It is a schematic diagram of the structure of a computer system of a terminal device or a server suitable for implementing an embodiment of the present invention. Figure 6 The terminal device or server shown is only an example and should not bring any limitation to the functions and scope of use of the embodiments of the present invention.

[0105] like Figure 6 As shown, the computer system 600 includes a central processing unit (CPU) 601, which can perform various appropriate actions and processes according to a program stored in a read-only memory (ROM) 602 or a program loaded from a storage part 608 into a random access memory (RAM) 603. In the RAM 603, various programs and data required for the operation of the system 600 are also stored. The CPU 601, the ROM 602, and the RAM 603 are connected to each other through a bus 604. An input / output (I / O) interface 605 is also connected to the bus 604.

[0106] The following components are connected to the I / O interface 605: an input section 606 including a keyboard, a mouse, etc.; an output section 607 including a cathode ray tube (CRT), a liquid crystal display (LCD), etc., and a speaker, etc.; a storage section 608 including a hard disk, etc.; and a communication section 609 including a network interface card such as a LAN card, a modem, etc. The communication section 609 performs communication processing via a network such as the Internet. A drive 610 is also connected to the I / O interface 605 as needed. A removable medium 611, such as a magnetic disk, an optical disk, a magneto-optical disk, a semiconductor memory, etc., is installed on the drive 610 as needed, so that a computer program read therefrom is installed into the storage section 608 as needed.

[0107] In particular, according to the embodiments disclosed in the present invention, the process described above with reference to the flowchart can be implemented as a computer software program. For example, the embodiments disclosed in the present invention include a computer program product, which includes a computer program carried on a computer-readable medium, and the computer program contains program code for executing the method shown in the flowchart. In such an embodiment, the computer program can be downloaded and installed from the network through the communication part 609, and / or installed from the removable medium 611. When the computer program is executed by the central processing unit (CPU) 601, the above-mentioned functions defined in the system of the present invention are executed.

[0108] It should be noted that the computer-readable medium shown in the present invention may be a computer-readable signal medium or a computer-readable storage medium or any combination of the two. The computer-readable storage medium may be, for example, but not limited to, an electrical, magnetic, optical, electromagnetic, infrared, or semiconductor system, device or device, or any combination of the above. More specific examples of computer-readable storage media may include, but are not limited to: an electrical connection with one or more wires, a portable computer disk, a hard disk, a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or flash memory), an optical fiber, a portable compact disk read-only memory (CD-ROM), an optical storage device, a magnetic storage device, or any suitable combination of the above. In the present invention, a computer-readable storage medium may be any tangible medium containing or storing a program that can be used by or in combination with an instruction execution system, device or device. In the present invention, a computer-readable signal medium may include a data signal propagated in a baseband or as part of a carrier wave, which carries a computer-readable program code. This propagated data signal may take a variety of forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination of the above. Computer-readable signal media may also be any computer-readable medium other than computer-readable storage media, which may send, propagate or transmit a program for use by or in conjunction with an instruction execution system, apparatus or device. The program code contained on the computer-readable medium may be transmitted using any appropriate medium, including but not limited to: wireless, wire, optical cable, RF, etc., or any suitable combination thereof.

[0109] The flow chart and block diagram in the accompanying drawings illustrate the possible architecture, function and operation of the system, method and computer program product according to various embodiments of the present invention. In this regard, each box in the flow chart or block diagram can represent a module, a program segment or a part of a code, and the module, program segment or a part of the code contains one or more executable instructions for realizing the specified logical function. It should also be noted that in some alternative implementations, the functions marked in the box can also occur in a sequence different from that marked in the accompanying drawings. For example, two boxes represented in succession can actually be executed substantially in parallel, and they can sometimes be executed in the opposite order, depending on the functions involved. It should also be noted that each box in the block diagram or flow chart, and the combination of the boxes in the block diagram or flow chart can be implemented with a dedicated hardware-based system that performs the specified function or operation, or can be implemented with a combination of dedicated hardware and computer instructions.

[0110] The units involved in the embodiments of the present invention may be implemented by software or hardware. The units described may also be set in a processor, for example, it may be described as: a processor includes: a first training module, a second training module and a third training module.

[0111] Among them, the names of these modules do not constitute a limitation on the modules themselves in certain cases. For example, the first training module can also be described as "a module used to obtain a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; and to use the first training sample to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model."

[0112] On the other hand, the present invention also provides a computer-readable medium, which may be included in the device described in the embodiment; or it may exist independently and not be assembled into the device. The computer-readable medium carries one or more programs, and when the one or more programs are executed by a device, the device includes: obtaining a first image sample and a bounding box of a target object annotated in the first image sample to generate a first training sample; using the first training sample to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model; using the first image segmentation model to annotate the image bounding box of a second image sample, and then using the annotated result and a general image segmentation model to train the image segmentation capability of the first image segmentation model to obtain a second image segmentation model; based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained to obtain a trained image segmentation model, and the scale of the image segmentation model is smaller than that of the general image segmentation model.

[0113] According to the technical solution of the embodiment of the present invention, the following advantages or beneficial effects are achieved: a first training sample is generated by acquiring a first image sample and a bounding box of a target object annotated in the first image sample; the first training sample is used to train the bounding box detection capability of an image segmentation model to obtain a first image segmentation model; the first image segmentation model is used to annotate the image bounding box of a second image sample, and the image segmentation capability of the first image segmentation model is trained using the annotated result and a general image segmentation model to obtain a second image segmentation model; based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained to obtain a trained image segmentation model; the scale of the image segmentation model is smaller than that of the general image segmentation model, and a lightweight image segmentation model training method is implemented. The bounding box detection capability of the image segmentation model and the general image segmentation model are used to automatically annotate the second image sample, thereby saving annotation time and cost; and the segmentation performance of the small-scale image segmentation model is guaranteed by overall optimization training of the bounding box detection capability and image segmentation capability of the model, thereby achieving a trained image segmentation model with good segmentation performance without occupying too many resources.

[0114] The specific implementation methods described herein do not constitute limitations on the scope of protection of the present invention. It should be understood by those skilled in the art that various modifications, combinations, sub-combinations and substitutions may occur depending on design requirements and other factors. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the present invention shall be included in the scope of protection of the present invention.

Claims

1. A training method for an image segmentation model, characterized in that: include: Acquire a first image sample and a bounding box of a target object marked in the first image sample to generate a first training sample; Using the first training sample to train the bounding box detection capability of the image segmentation model to obtain a first image segmentation model; Using the first image segmentation model to annotate the image bounding box of the second image sample, and then using the annotated result and the general image segmentation model to train the image segmentation capability of the first image segmentation model to obtain a second image segmentation model; Based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than that of the general graphic segmentation model.

2. The method according to claim 1, characterized in that Using the first image segmentation model to annotate the image bounding box of the second image sample, and then using the annotated result and the general image segmentation model to train the image segmentation capability of the first image segmentation model, including: Inputting a second image sample into the first image segmentation model, and obtaining a bounding box of a target object in the second image sample and a corresponding segmented image output by the first image segmentation model; Inputting the second image sample and the bounding box of the target object in the second image sample into a general image segmentation model to obtain a standard segmentation image output by the general image segmentation model; The image segmentation capability of the first image segmentation model is trained by calculating the minimum loss function between the standard segmentation image and the segmentation image.

3. The method according to claim 1, characterized in that: Based on the bounding box detection capability and image segmentation capability of the second image segmentation model, the second image segmentation model is overall optimized and trained, including: Inputting the third image sample into the second image segmentation model, and using the bounding box detection capability and image segmentation capability of the second image segmentation model to perform bounding box annotation and segmentation image annotation on the target object in the third image sample; Screening the bounding box of the target object marked in the third image sample and the corresponding segmented image, and generating a second training sample according to the screening result; The second image segmentation model is trained using the second training sample.

4. The method according to claim 1, characterized in that: Before obtaining the first image sample and the bounding box of the target object marked in the first image sample, the method further includes: The bounding box of the target object in the first image sample is annotated according to the annotation rule of the general image segmentation model, wherein the annotation rule is to use a specified mark at a specified position to characterize the position of the target object in the first image sample.

5. A method for image segmentation, characterized in that: include: Obtain the image to be segmented; The image to be segmented is input into a pre-constructed image segmentation model, and a bounding box and a segmented image of the target object are obtained as output by the image segmentation model; wherein the bounding box indicates the position of the target object in the image to be segmented, the segmented image is segmented from the image to be segmented and contains the target object, and the image segmentation model is constructed according to any one of the training methods of claims 1 to 4.

6. The method according to claim 5, characterized in that Inputting the image to be segmented into a pre-built image segmentation model, and obtaining a bounding box and a segmented image of the target object output by the image segmentation model, comprising: Inputting the image to be segmented into the image segmentation model, performing feature extraction on the image to be segmented through the backbone network of the image segmentation model, and obtaining image features of the image to be segmented; Performing bounding box detection of the target object on the image features through the detection network of the image segmentation model, and obtaining the bounding box of the target object output by the image segmentation model; The image features are subjected to image segmentation of a target object through a segmentation network of the image segmentation model, and a segmented image of the target object output by the image segmentation model is obtained.

7. An image annotation method, characterized in that: include: Obtain the image to be annotated; The image to be annotated is input into a pre-constructed image segmentation model, and the image to be annotated is annotated according to a bounding box of a target object and a corresponding segmented image output by the image segmentation model, wherein the image segmentation model is constructed according to any one of the training methods in claims 1 to 4.

8. A training device for an image segmentation model, characterized in that: include: A first training module, configured to obtain a first image sample and a bounding box of a target object marked in the first image sample, and generate a first training sample; Using the first training sample to train the bounding box detection capability of the image segmentation model to obtain a first image segmentation model; A second training module is used to use the first image segmentation model to annotate the image bounding box of the second image sample, and then use the annotation result and the general image segmentation model to train the image segmentation ability of the first image segmentation model to obtain a second image segmentation model; The third training module is used to perform overall optimization training on the second image segmentation model based on the bounding box detection capability and image segmentation capability of the second image segmentation model to obtain a trained image segmentation model, wherein the scale of the image segmentation model is smaller than that of the general graphic segmentation model.

9. A mobile electronic device terminal, characterized in that: include: one or more processors; a storage device for storing one or more programs, When the one or more programs are executed by the one or more processors, the one or more processors implement the method according to any one of claims 1 to 4.

10. A computer readable medium having a computer program stored thereon, characterized in that: When the program is executed by a processor, the method according to any one of claims 1 to 4 is implemented.

11. A computer program product, comprising a computer program, characterized in that When the computer program is executed by a processor, the method according to any one of claims 1 to 4 is implemented.

Images