Safety helmet data generation method and device based on generative model

Abstract

The invention relates to a safety helmet data generation method and device based on a generative model. The method comprises the following steps: acquiring safety helmet data, marking a text prompt dialog box and a safety helmet bounding box, and segmenting and marking the safety helmet data. And training and finely adjusting the Mobile SAM lightweight segmentation model through the segmented and labeled safety helmet data, and finely adjusting the CLIP encoder through a text prompt dialog box. And calling the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the safety helmet data in the safety helmet bounding box to obtain mask information of the safety helmet. And based on the text prompt dialog box, the safety helmet bounding box and the mask information, performing fine tuning and constraint on the Gligen model to obtain the Gligen model for generating safety helmet data. And taking the current safety helmet image as the input of the Gligen model, and calling the Gligen model to output safety helmet data corresponding to the current safety helmet image.

Description

Technical Field

[0001] The present invention relates to the technical field of safety management and intelligent monitoring, and in particular to a method and device for generating safety helmet data based on a generation model. Background Art

[0002] In the current field of industrial safety management and intelligent monitoring, helmet recognition technology is an important part of ensuring the safety of workers on the job site. However, building an efficient and accurate helmet detection system faces a major challenge, namely the scarcity of high-quality and diverse training data sets.

[0003] At present, common image generation technologies include generative adversarial networks (GANs) and variational autoencoders (VAEs). Among them, GANs can generate very realistic images, especially good at capturing the details and realism of images. However, its model training process is unstable, and it requires careful balance between the generator and the discriminator, which is prone to mode collapse (generated images lack diversity) and overfitting. In addition, the training of GANs usually requires a lot of computing resources and time. VAEs provide a mathematical framework for generating images, allowing the generation process to be explained, and the training is relatively stable. However, the quality of the generated images is usually not as good as that of the most advanced GANs, and may appear blurry or lack details. VAEs tend to generate overly smooth outputs because they maximize the lower bound of the data rather than directly maximize the likelihood. However, traditional GANs, VAEs, and other image generation algorithms generate images with poor details and textures, are less stable to train, and are more difficult to generate more diverse images, resulting in difficulty in generating real and effective helmet data. Summary of the invention

[0004] Based on this, it is necessary to provide a method and device for generating helmet data based on a generative model, which can improve the authenticity and effectiveness of helmet data generation, in response to the above technical problems.

[0005] The present invention provides a method for generating helmet data based on a generation model, the method comprising:

[0006] Acquire safety helmet data, and mark a text prompt dialog box and a safety helmet boundary box based on the safety helmet data to segment and annotate the safety helmet data;

[0007] The Mobile SAM lightweight segmentation model is trained and fine-tuned by segmenting and annotating the helmet data, and the CLIP encoder is fine-tuned by the text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model;

[0008] Calling the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data;

[0009] Based on the text prompt dialog box, the helmet bounding box and the mask information, fine-tuning and constraining the Gligen model to guide the Gligen model to generate edge data of the helmet, thereby obtaining a Gligen model for generating helmet data;

[0010] The current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image, wherein the helmet data includes the type, color and size of the helmet.

[0011] In one embodiment, the step of acquiring the helmet data and marking a text prompt dialog box and a helmet boundary box based on the helmet data to segment and annotate the helmet data includes:

[0012] Annotating the hard hat in each hard hat sample image in the hard hat data with the hard hat boundary box, and adding a text prompt dialog box to the hard hat boundary box;

[0013] The Labelme annotation tool is used to segment and annotate the helmet samples in the helmet boundary box to obtain the segmented and annotated helmet data.

[0014] In one embodiment, the training and fine-tuning of the Mobile SAM lightweight segmentation model by the segmented and labeled helmet data, and the fine-tuning of the CLIP encoder by the text prompt dialog box to obtain the trained Mobile SAM lightweight segmentation model include:

[0015] The segmented and labeled helmet data is used as the input of the Mobile SAM lightweight segmentation model, and the Mobile SAM lightweight segmentation model is trained and fine-tuned to enhance the segmentation function of the Mobile SAM lightweight segmentation model for helmets;

[0016] Using the text prompt dialog box corresponding to the segmented and marked helmet data as input of the CLIP encoder, fine-tuning the CLIP encoder, and calling the fine-tuned CLIP encoder to encode the text modality data;

[0017] The text modal data includes the text prompt dialog box, the helmet boundary box and the mask information, and the encoded text modal data is used as the input of the Gligen model.

[0018] In one embodiment, calling the trained Mobile SAM lightweight segmentation model to segment the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data includes:

[0019] The trained Mobile SAM lightweight segmentation model is called to crop the helmet image area based on the helmet bounding box of the segmented and labeled helmet data, and the helmet image area is segmented to obtain mask information corresponding to the helmet data;

[0020] The mask information is used to record the area information where the safety helmet is located contained in the safety helmet boundary box.

[0021] In one embodiment, fine-tuning and constraining the Gligen model based on the text prompt dialog box, the helmet bounding box and the mask information to guide the Gligen model to generate edge data of the helmet, and obtaining the Gligen model for generating helmet data includes:

[0022] Calling the fine-tuned CLIP encoder to encode the text prompt dialog box and the helmet bounding box, and using the encoded text prompt dialog box and the helmet bounding box as the Gligen model for training and fine-tuning; and

[0023] During the Gligen model training and fine-tuning process, the Gligen model is constrained by the encoded mask information to guide the Gligen model to generate edge data of the helmet.

[0024] In one embodiment, the step of using the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image includes:

[0025] Acquire the current hard hat image, and mark a hard hat bounding box of the current hard hat in the current hard hat image;

[0026] The current helmet within the helmet boundary box is segmented using the trained Mobile SAM lightweight segmentation model according to a preset text prompt dialog box to obtain mask information corresponding to the current helmet.

[0027] In one embodiment, the step of using the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image further includes:

[0028] The mask information, the helmet bounding box and the preset text prompt dialog box corresponding to the current helmet are used as inputs of the Gligen model after fine-tuning and constraint, so as to call the Gligen model to output the helmet data corresponding to the current helmet.

[0029] The present invention also provides a device for generating helmet data based on a generation model, the device comprising:

[0030] A segmentation and annotation module is used to obtain the helmet data and mark a text prompt dialog box and a helmet boundary box based on the helmet data to segment and annotate the helmet data;

[0031] The first model training module is used to train and fine-tune the Mobile SAM lightweight segmentation model through the segmented and labeled helmet data, and fine-tune the CLIP encoder through the text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model;

[0032] A mask information generating module, used for calling the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data;

[0033] A second model training module is used to fine-tune and constrain the Gligen model based on the text prompt dialog box, the helmet bounding box and the mask information to guide the Gligen model to generate edge data of the helmet, thereby obtaining a Gligen model for generating helmet data;

[0034] The helmet data generation module is used to use the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image, wherein the helmet data includes the type, color and size of the helmet.

[0035] The present invention also provides an electronic device, comprising a memory and a processor, wherein the memory stores a computer program, and when the processor executes the computer program, the method for generating helmet data based on a generation model as described above is implemented.

[0036] The present invention also provides a computer storage medium storing a computer program, wherein when the computer program is executed by a processor, the method for generating helmet data based on a generation model as described above is implemented.

[0037] The above-mentioned method and device for generating helmet data based on the generative model obtains the helmet data and marks the text prompt dialog box and the helmet bounding box based on the helmet data to segment and annotate the helmet data. The Mobile SAM lightweight segmentation model is then trained and fine-tuned through the segmented and annotated helmet data, and the CLIP encoder is fine-tuned through the text prompt dialog box to obtain the trained Mobile SAM lightweight segmentation model. Subsequently, the trained Mobile SAM lightweight segmentation model is called to segment the helmet data within the helmet bounding box to obtain the mask information corresponding to the helmet data. Based on the text prompt dialog box, the helmet bounding box and the mask information, the Gligen model is fine-tuned and constrained to guide the Gligen model to generate the edge data of the helmet, and the Gligen model for generating helmet data is obtained. Finally, the current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image. This method introduces Gligen, an advanced artificial intelligence image generation technology. Compared with traditional generative adversarial networks (GANs), the Gligen model learns data distribution by gradually adding and removing noise in images. At the same time, it introduces two modal data, text and helmet bounding boxes, which can generate more realistic and detailed images. It can also better control the diversity and consistency of generated images, which improves the authenticity and effectiveness of helmet data generation to a certain extent. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] In order to more clearly illustrate the technical solutions in the present invention or the prior art, the following briefly introduces the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative work.

[0039] Figure 1 One of the flow charts of the method for generating helmet data based on the generation model provided by the present invention;

[0040] Figure 2 A schematic diagram of a safety helmet data generation process of a safety helmet data generation method based on a generation model in a specific embodiment provided by the present invention;

[0041] Figure 3 The second flowchart of the method for generating helmet data based on the generation model provided by the present invention;

[0042] Figure 4 The third flowchart of the method for generating helmet data based on the generation model provided by the present invention;

[0043] Figure 5 The fourth flowchart of the method for generating helmet data based on the generation model provided by the present invention;

[0044] Figure 6 The fifth flowchart of the method for generating helmet data based on the generation model provided by the present invention;

[0045] Figure 7 A schematic diagram of the structure of a device for generating helmet data based on a generation model provided by the present invention;

[0046] Figure 8 This is a diagram of the internal structure of the electronic device provided by the present invention. DETAILED DESCRIPTION

[0047] In order to make the purpose, technical solution and advantages of the embodiments of the present invention clearer, the technical solution in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0048] Combine the following Figures 1 to 8 The present invention describes a method and device for generating helmet data based on a generation model.

[0049] like Figure 1 As shown, in one embodiment, a method for generating helmet data based on a generation model comprises the following steps:

[0050] Step S110 , obtaining helmet data, and marking a text prompt dialog box and a helmet boundary box based on the helmet data to segment and annotate the helmet data.

[0051] Specifically, the server obtains real helmet data, which includes helmets of various styles, colors, and sizes, and marks appropriate text prompt dialog boxes (Prompt) and helmet bounding boxes (Bounding box) based on the collected helmet data, and then segments and annotates the helmets in the bounding box.

[0052] Combination Figure 2As shown, in a specific embodiment, the method for generating helmet data based on a generative model provided by the present invention first needs to collect some real helmet data, including helmets of various styles, colors, and sizes, and mark appropriate text prompts, helmet bounding boxes, and segment and annotate the helmets for the collected helmet data. For each helmet in a sample image, a bounding box is annotated, and an appropriate text prompt, i.e., a text prompt dialog box, is added to the helmet in the bounding box, such as "a black helmet". The helmet segmentation is mainly annotated using the Labelme annotation tool.

[0053] In this embodiment, all text prompts can be in a custom format, for example, "one+color+safety helmet" or "one+style+safety helmet". Each safety helmet Boundingbox has a corresponding text prompt, and the safety helmets in each Boundingbox area will be image segmented and labeled.

[0054] Step S120 , training and fine-tuning the Mobile SAM lightweight segmentation model by segmenting and annotating the helmet data, and fine-tuning the CLIP encoder through a text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model.

[0055] Specifically, the server trains and fine-tunes the Mobile SAM lightweight segmentation model by segmenting and annotating the helmet data to enhance the Mobile SAM lightweight segmentation model's ability to segment helmets, and fine-tunes the CLIP encoder through the text prompt dialog box. All subsequent text modal data must use the fine-tuned CLIP encoder as input to the Gligen model.

[0056] Combination Figure 2 As shown, in a specific embodiment, the method for generating helmet data based on a generative model provided by the present invention selects a Mobile SAM lightweight segmentation model, and uses segmented and labeled helmet data for fine-tuning to enhance the model's ability to segment helmets, thereby improving the accuracy of helmet segmentation. Subsequent Gligen model training needs to rely on the corresponding helmet segmentation result Mask, thereby enhancing the generation and reconstruction of the Mask area by the Gligen model, rather than generating and reconstructing the entire Boundingbox area.

[0057] In the present embodiment, it is also necessary to use the marked helmet image text prompt to fine-tune the CLIP encoder, and all subsequent text modal data need to be encoded using the fine-tuned CLIP encoder to serve as the input of the Gligen model. Wherein, the CLIP encoder is a text encoder for encoding text data into feature vectors, which is the input format of the Gligen model. The purpose of fine-tuning the CLIP encoder is to allow the CLIP encoder to learn the potential feature relationship between the constructed helmet data and the corresponding text prompt dialog box, so that the subsequent Gligen model can fuse the two. In addition, the fine-tuned CLIP encoder can associate the data of the two different modes of the helmet image and the text prompt dialog box, so as to facilitate the generation of helmet data.

[0058] Step S130, calling the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data.

[0059] Specifically, the server calls the trained Mobile SAM lightweight segmentation model to segment the helmet data within the helmet bounding box to obtain the mask information required for training, that is, Mask information, and the Mask information records the area information where the helmet is located contained in the helmet bounding box.

[0060] Combination Figure 2 As shown, in a specific embodiment, the method for generating helmet data based on a generative model provided by the present invention uses a trained segmentation model Mobile SAM lightweight segmentation model to crop the image area containing the helmet based on the Boundingbox information of the labeled helmet to segment it, and then obtain the Mask information required for training. The Mask information records the area where the helmet is located contained in the Bounding box.

[0061] Step S140, based on the text prompt dialog box, the helmet bounding box and the mask information, fine-tune and constrain the Gligen model to guide the Gligen model to generate edge data of the helmet, and obtain the Gligen model for generating helmet data.

[0062] Specifically, the server fine-tunes the Gligen model based on the text prompt dialog box and the helmet bounding box obtained above, and constrains the Gligen model based on the mask information to guide the Gligen model to generate edge data of the helmet, thereby obtaining the Gligen model for generating helmet data.

[0063] Combination Figure 2As shown, in a specific embodiment, the method for generating helmet data based on a generative model provided by the present invention, after obtaining the Mask information required for training, uses the marked helmet image text Prompt and Bounding box data to fine-tune the Gligen model, and at the same time uses the generated helmet Mask information to constrain the Gligen model, guiding the Gligen model to focus on the generation of edge details of the helmet, so that the Gligen model has the ability to generate helmet data.

[0064] Step S150, taking the current helmet image as the input of the Gligen model, so as to call the Gligen model to output the helmet data corresponding to the current helmet image, where the helmet data includes the type, color and size of the helmet.

[0065] Specifically, the server uses the current helmet image to be processed as the input of the Gligen model for generating helmet data obtained in step S140, and calls the Gligen model to output the helmet data corresponding to the current helmet image, that is, the style type, color and size of the current helmet data.

[0066] Combination Figure 2 As shown, in a specific embodiment, the method for generating helmet data based on a generative model provided by the present invention, when using the Gligen model fine-tuned with the labeled helmet data for inference, can randomly select a picture containing a helmet as input, and mark the Bounding box area of ​​the helmet in the input image. For example, the helmet in the input image is red, and it is hoped that the red helmet in the original image will be replaced with a black helmet. Then the text Prompt can be set to "a black helmet", and at the same time, the trained Mobile SAM lightweight segmentation model is used to segment the helmet in the Bounding box area to obtain the corresponding Mask information, and finally the Bounding box information and the Prompt text description are combined to input the Gligen model network to generate the required helmet image.

[0067] The above-mentioned method for generating helmet data based on the generative model obtains the helmet data and marks the text prompt dialog box and the helmet bounding box based on the helmet data to segment and annotate the helmet data. The Mobile SAM lightweight segmentation model is then trained and fine-tuned through the segmented and annotated helmet data, and the CLIP encoder is fine-tuned through the text prompt dialog box to obtain the trained Mobile SAM lightweight segmentation model. Subsequently, the trained Mobile SAM lightweight segmentation model is called to segment the helmet data within the helmet bounding box to obtain the mask information corresponding to the helmet data. Based on the text prompt dialog box, the helmet bounding box and the mask information, the Gligen model is fine-tuned and constrained to guide the Gligen model to generate the edge data of the helmet, and the Gligen model for generating helmet data is obtained. Finally, the current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image. This method introduces Gligen, an advanced artificial intelligence image generation technology. Compared with traditional generative adversarial networks (GANs), the Gligen model learns data distribution by gradually adding and removing noise in images. At the same time, it introduces two modal data, text and helmet bounding boxes, which can generate more realistic and detailed images. It can also better control the diversity and consistency of generated images, which improves the authenticity and effectiveness of helmet data generation to a certain extent.

[0068] like Figure 3 As shown, in one embodiment, the method for generating helmet data based on a generative model provided by the present invention obtains helmet data, and marks a text prompt dialog box and a helmet boundary box based on the helmet data to segment and annotate the helmet data, specifically comprising the following steps:

[0069] Step S112: annotate a safety helmet boundary box for the safety helmet in each safety helmet sample image in the safety helmet data, and add a text prompt dialog box to the safety helmet boundary box.

[0070] Step S114, using the Labelme annotation tool to segment and annotate the helmet samples in the helmet boundary box to obtain segmented and annotated helmet data.

[0071] like Figure 4 As shown, in one embodiment, the method for generating helmet data based on a generative model provided by the present invention trains and fine-tunes the Mobile SAM lightweight segmentation model by segmenting and annotating the helmet data, and fine-tunes the CLIP encoder through a text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model, which specifically includes the following steps:

[0072] Step S122, using the segmented and labeled helmet data as the input of the Mobile SAM lightweight segmentation model, training and fine-tuning the Mobile SAM lightweight segmentation model to enhance the helmet segmentation function of the Mobile SAM lightweight segmentation model.

[0073] Step S124, using the text prompt dialog box corresponding to the segmented and annotated helmet data as the input of the CLIP encoder, fine-tuning the CLIP encoder, and calling the fine-tuned CLIP encoder to encode the text modality data.

[0074] The text modal data includes a text prompt dialog box, a helmet bounding box, and mask information. The encoded text modal data is used as the input of the Gligen model.

[0075] like Figure 5 As shown, in one embodiment, the method for generating helmet data based on the generative model provided by the present invention fine-tunes and constrains the Gligen model based on the text prompt dialog box, the helmet bounding box and the mask information to guide the Gligen model to generate the edge data of the helmet, and obtains the Gligen model for generating the helmet data, which specifically includes the following steps:

[0076] Step S142, calling the fine-tuned CLIP encoder to encode the text prompt dialog box and the helmet bounding box, and using the encoded text prompt dialog box and the helmet bounding box as the Gligen model for training and fine-tuning.

[0077] Step S144, during the Gligen model training and fine-tuning process, the Gligen model is constrained by the encoded mask information to guide the Gligen model to generate edge data of the helmet.

[0078] like Figure 6 As shown, in one embodiment, the method for generating helmet data based on a generative model provided by the present invention takes the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image, specifically comprising the following steps:

[0079] Step S152, obtaining the current helmet image, and marking the helmet boundary box of the current helmet in the current helmet image.

[0080] Step S154, using the trained Mobile SAM lightweight segmentation model to segment the current helmet within the helmet boundary box according to a preset text prompt dialog box to obtain mask information corresponding to the current helmet.

[0081] Step S156, using the mask information corresponding to the current helmet, the helmet bounding box and the preset text prompt dialog box as inputs of the fine-tuned and constrained Gligen model, so as to call the Gligen model to output the helmet data corresponding to the current helmet.

[0082] The following is a description of a safety helmet data generating device based on a generative model provided by the present invention. The safety helmet data generating device based on a generative model described below and the safety helmet data generating method based on a generative model described above can correspond to each other.

[0083] like Figure 7 As shown, in one embodiment, a device for generating helmet data based on a generative model includes a segmentation and annotation module 710 , a first model training module 720 , a mask information generation module 730 , a second model training module 740 and a helmet data generation module 750 .

[0084] The segmentation and annotation module 710 is used to obtain the helmet data, and mark the text prompt dialog box and the helmet boundary box based on the helmet data to segment and annotate the helmet data.

[0085] The first model training module 720 is used to train and fine-tune the Mobile SAM lightweight segmentation model through the segmented and labeled helmet data, and fine-tune the CLIP encoder through a text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model.

[0086] The mask information generating module 730 is used to call the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data.

[0087] The second model training module 740 is used to fine-tune and constrain the Gligen model based on the text prompt dialog box, the helmet bounding box and the mask information to guide the Gligen model to generate edge data of the helmet, and obtain the Gligen model for generating helmet data.

[0088] The helmet data generation module 750 is used to use the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image. The helmet data includes the type, color and size of the helmet.

[0089] In this embodiment, in the safety helmet data generation device based on the generation model provided by the present invention, the segmentation and annotation module 710 is specifically used for:

[0090] Annotate the hard hat bounding box for each hard hat sample image in the hard hat data, and add a text prompt dialog box to the hard hat bounding box.

[0091] The Labelme annotation tool is used to segment and annotate the helmet samples in the helmet boundary box to obtain the segmented and annotated helmet data.

[0092] In this embodiment, in the safety helmet data generating device based on the generation model provided by the present invention, the first model training module 720 is specifically used for:

[0093] The segmented and labeled helmet data is used as the input of the Mobile SAM lightweight segmentation model, and the MobileSAM lightweight segmentation model is trained and fine-tuned to enhance the segmentation function of the Mobile SAM lightweight segmentation model for helmets.

[0094] The text prompt dialog box corresponding to the segmented and annotated helmet data is used as the input of the CLIP encoder, the CLIP encoder is fine-tuned, and the fine-tuned CLIP encoder is called to encode the text modal data.

[0095] The text modal data includes a text prompt dialog box, a helmet bounding box, and mask information. The encoded text modal data is used as the input of the Gligen model.

[0096] In this embodiment, the safety helmet data generating device based on the generation model provided by the present invention, the mask information generating module 730 is specifically used for:

[0097] The trained Mobile SAM lightweight segmentation model is called to crop the helmet image area based on the helmet bounding box of the segmented and labeled helmet data, and the helmet image area is segmented to obtain the mask information corresponding to the helmet data.

[0098] The mask information is used to record the area information where the helmet is located within the helmet boundary box.

[0099] In this embodiment, in the safety helmet data generating device based on the generation model provided by the present invention, the second model training module 740 is specifically used for:

[0100] The fine-tuned CLIP encoder is called to encode the text prompt dialog box and the helmet bounding box, and the encoded text prompt dialog box and the helmet bounding box are used as the Gligen model for training and fine-tuning.

[0101] During the Gligen model training and fine-tuning process, the Gligen model is constrained by the encoded mask information to guide the Gligen model to generate edge data of the helmet.

[0102] In this embodiment, the helmet data generating device based on the generation model provided by the present invention, the helmet data generating module 750 is specifically used for:

[0103] Get the current hard hat image and mark the hard hat bounding box of the current hard hat in the current hard hat image.

[0104] The trained Mobile SAM lightweight segmentation model is used to segment the current helmet within the helmet bounding box according to the preset text prompt dialog box to obtain the mask information corresponding to the current helmet.

[0105] In this embodiment, the helmet data generating device based on the generation model provided by the present invention, the helmet data generating module 750 is specifically used for:

[0106] The mask information, the helmet bounding box and the preset text prompt dialog box corresponding to the current helmet are used as the input of the fine-tuned and constrained Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet.

[0107] Figure 8 An example of a physical structure diagram of an electronic device is shown. The electronic device may be a smart terminal, and its internal structure diagram may be as follows: Figure 8 As shown. The electronic device includes a processor, an internal memory and a network interface connected through a system bus. The processor of the electronic device is used to provide computing and control capabilities. The memory of the electronic device includes a non-volatile storage medium and an internal memory. The non-volatile storage medium stores an operating system and a computer program. The internal memory provides an environment for the operation of the operating system and the computer program in the non-volatile storage medium. The network interface of the electronic device is used to communicate with an external terminal through a network connection. When the computer program is executed by the processor, a method for generating helmet data based on a generation model is implemented, and the method includes:

[0108] Acquire the hard hat data, and mark the text prompt dialog box and the hard hat bounding box based on the hard hat data to segment and annotate the hard hat data;

[0109] The Mobile SAM lightweight segmentation model is trained and fine-tuned by segmenting and annotating the helmet data, and the CLIP encoder is fine-tuned through the text prompt dialog box to obtain the trained Mobile SAM lightweight segmentation model.

[0110] Call the trained Mobile SAM lightweight segmentation model to segment the helmet data within the helmet bounding box to obtain the mask information corresponding to the helmet data;

[0111] Based on the text prompt dialog box, the helmet bounding box and the mask information, the Gligen model is fine-tuned and constrained to guide the Gligen model to generate the edge data of the helmet, and the Gligen model for generating the helmet data is obtained;

[0112] The current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image. The helmet data includes the type, color and size of the helmet.

[0113] Those skilled in the art will understand that Figure 8 The structure shown in the figure is only a block diagram of a part of the structure related to the scheme of the present invention, and does not constitute a limitation on the electronic device to which the scheme of the present invention is applied. The specific electronic device may include more or fewer components than those shown in the figure, or combine certain components, or have a different arrangement of components.

[0114] On the other hand, the present invention also provides a computer storage medium storing a computer program, which, when executed by a processor, implements a method for generating helmet data based on a generation model, the method comprising:

[0115] Acquire the hard hat data, and mark the text prompt dialog box and the hard hat bounding box based on the hard hat data to segment and annotate the hard hat data;

[0116] The Mobile SAM lightweight segmentation model is trained and fine-tuned by segmenting and annotating the helmet data, and the CLIP encoder is fine-tuned through the text prompt dialog box to obtain the trained Mobile SAM lightweight segmentation model.

[0117] Call the trained Mobile SAM lightweight segmentation model to segment the helmet data within the helmet bounding box to obtain the mask information corresponding to the helmet data;

[0118] Based on the text prompt dialog box, the helmet bounding box and the mask information, the Gligen model is fine-tuned and constrained to guide the Gligen model to generate the edge data of the helmet, and the Gligen model for generating the helmet data is obtained;

[0119] The current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image. The helmet data includes the type, color and size of the helmet.

[0120] In another aspect, a computer program product or a computer program is provided, the computer program product or the computer program comprising computer instructions, the computer instructions being stored in a computer-readable storage medium. A processor of an electronic device reads the computer instructions from the computer-readable storage medium, and when the processor executes the computer instructions, a method for generating helmet data based on a generative model is implemented, the method comprising:

[0121] Acquire the hard hat data, and mark the text prompt dialog box and the hard hat bounding box based on the hard hat data to segment and annotate the hard hat data;

[0122] The Mobile SAM lightweight segmentation model is trained and fine-tuned by segmenting and annotating the helmet data, and the CLIP encoder is fine-tuned through the text prompt dialog box to obtain the trained Mobile SAM lightweight segmentation model.

[0123] Call the trained Mobile SAM lightweight segmentation model to segment the helmet data within the helmet bounding box to obtain the mask information corresponding to the helmet data;

[0124] Based on the text prompt dialog box, the helmet bounding box and the mask information, the Gligen model is fine-tuned and constrained to guide the Gligen model to generate the edge data of the helmet, and the Gligen model for generating the helmet data is obtained;

[0125] The current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image. The helmet data includes the type, color and size of the helmet.

[0126] Those of ordinary skill in the art can understand that all or part of the processes in the above-mentioned embodiments can be implemented by instructing the relevant hardware through a computer program, and the computer program can be stored in a non-volatile computer-readable storage medium. When the computer program is executed, it can include the processes of the embodiments of the above-mentioned methods. Among them, any reference to memory, storage, database or other media used in the embodiments provided by the present invention may include non-volatile and / or volatile memory. Non-volatile memory may include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM) or flash memory. Volatile memory may include random access memory (RAM) or external cache memory.

[0127] By way of illustration and not limitation, RAM is available in many forms, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), double data rate SDRAM (DDRSDRAM), enhanced SDRAM (ESDRAM), synchronous link (Synchlink) DRAM (SLDRAM), memory bus (Rambus) direct RAM (RDRAM), direct memory bus dynamic RAM (DRDRAM), and memory bus dynamic RAM (RDRAM), etc.

[0128] The technical features of the above embodiments may be arbitrarily combined. To make the description concise, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0129] The above-mentioned embodiments only express several implementation methods of the present invention, and the description thereof is relatively specific and detailed, but it cannot be understood as limiting the scope of the patent of the present invention. It should be pointed out that, for ordinary technicians in this field, several variations and improvements can be made without departing from the concept of the present invention, which all belong to the protection scope of the present invention. Therefore, the protection scope of the patent of the present invention shall be subject to the attached claims.

Claims

1. A method for generating helmet data based on a generative model, characterized in that: The method comprises: Acquire safety helmet data, and mark a text prompt dialog box and a safety helmet boundary box based on the safety helmet data to segment and annotate the safety helmet data; The Mobile SAM lightweight segmentation model is trained and fine-tuned by segmenting and annotating the helmet data, and the CLIP encoder is fine-tuned by the text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model; Calling the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data; Based on the text prompt dialog box, the helmet bounding box and the mask information, fine-tuning and constraining the Gligen model to guide the Gligen model to generate edge data of the helmet, thereby obtaining a Gligen model for generating helmet data; The current helmet image is used as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image, wherein the helmet data includes the type, color and size of the helmet.

2. The method for generating helmet data based on a generative model according to claim 1, characterized in that: The step of acquiring the helmet data and marking a text prompt dialog box and a helmet boundary box based on the helmet data to segment and annotate the helmet data includes: Annotating the hard hat in each hard hat sample image in the hard hat data with the hard hat boundary box, and adding a text prompt dialog box to the hard hat boundary box; The Labelme annotation tool is used to segment and annotate the helmet samples in the helmet boundary box to obtain the segmented and annotated helmet data.

3. The method for generating helmet data based on a generative model according to claim 2, characterized in that: The Mobile SAM lightweight segmentation model is trained and fine-tuned by the segmented and labeled helmet data, and the CLIP encoder is fine-tuned by the text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model, including: The segmented and labeled helmet data is used as the input of the Mobile SAM lightweight segmentation model, and the Mobile SAM lightweight segmentation model is trained and fine-tuned to enhance the segmentation function of the Mobile SAM lightweight segmentation model for helmets; Using the text prompt dialog box corresponding to the segmented and marked helmet data as input of the CLIP encoder, fine-tuning the CLIP encoder, and calling the fine-tuned CLIP encoder to encode the text modality data; The text modal data includes the text prompt dialog box, the helmet boundary box and the mask information, and the encoded text modal data is used as the input of the Gligen model.

4. The method for generating helmet data based on a generative model according to claim 3, characterized in that: The calling of the trained Mobile SAM lightweight segmentation model to segment the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data includes: The trained Mobile SAM lightweight segmentation model is called to crop the helmet image area based on the helmet bounding box of the segmented and labeled helmet data, and the helmet image area is segmented to obtain mask information corresponding to the helmet data; The mask information is used to record the area information where the safety helmet is located contained in the safety helmet boundary box.

5. The method for generating helmet data based on a generative model according to claim 4, characterized in that: The method further comprises: fine-tuning and constraining the Gligen model based on the text prompt dialog box, the helmet bounding box and the mask information to guide the Gligen model to generate edge data of the helmet, and obtaining the Gligen model for generating helmet data, including: Calling the fine-tuned CLIP encoder to encode the text prompt dialog box and the helmet bounding box, and using the encoded text prompt dialog box and the helmet bounding box as the Gligen model for training and fine-tuning; and During the Gligen model training and fine-tuning process, the Gligen model is constrained by the encoded mask information to guide the Gligen model to generate edge data of the helmet.

6. The method for generating helmet data based on a generative model according to claim 5, characterized in that: The step of taking the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image includes: Acquire the current hard hat image, and mark a hard hat bounding box of the current hard hat in the current hard hat image; The current helmet within the helmet boundary box is segmented using the trained Mobile SAM lightweight segmentation model according to a preset text prompt dialog box to obtain mask information corresponding to the current helmet.

7. The method for generating helmet data based on a generative model according to claim 6, characterized in that: The method of using the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image further includes: The mask information, the helmet bounding box and the preset text prompt dialog box corresponding to the current helmet are used as inputs of the Gligen model after fine-tuning and constraint, so as to call the Gligen model to output the helmet data corresponding to the current helmet.

8. A device for generating helmet data based on a generative model, characterized in that: The device comprises: A segmentation and annotation module is used to obtain the helmet data and mark a text prompt dialog box and a helmet boundary box based on the helmet data to segment and annotate the helmet data; The first model training module is used to train and fine-tune the Mobile SAM lightweight segmentation model through the segmented and labeled helmet data, and fine-tune the CLIP encoder through the text prompt dialog box to obtain a trained Mobile SAM lightweight segmentation model; A mask information generating module, used for calling the trained Mobile SAM lightweight segmentation model to perform segmentation processing on the helmet data within the helmet boundary box to obtain mask information corresponding to the helmet data; A second model training module is used to fine-tune and constrain the Gligen model based on the text prompt dialog box, the helmet bounding box and the mask information to guide the Gligen model to generate edge data of the helmet, thereby obtaining a Gligen model for generating helmet data; The helmet data generation module is used to use the current helmet image as the input of the Gligen model to call the Gligen model to output the helmet data corresponding to the current helmet image, wherein the helmet data includes the type, color and size of the helmet.

9. An electronic device comprising a memory and a processor, wherein the memory stores a computer program, wherein: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 7 are implemented.

10. A computer storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

Images