Mars image augmentation method, computing device and storage medium based on stable diffusion model

Abstract

The present invention provides a Mars image augmentation method, computing device, and storage medium based on a stable diffusion model, and relates to the field of image processing technology. In the present invention, a training sample set is obtained, the training sample set including multiple Mars images and explanatory information corresponding to each of the multiple Mars images; the stable diffusion model is trained according to the training sample set to obtain a trained stable diffusion model; based on constraints, the trained stable diffusion model is conditionally trained to obtain a Control Net model; and an augmented image is obtained through the trained stable diffusion model and the Control Net model according to input feature information, wherein the feature information includes feature text and / or feature image. The method provided by the present invention implements image augmentation processing of Mars images based on a diffusion model, and can quickly obtain a large number of high-quality Mars images that meet image diversity requirements, thereby meeting the user's Mars geological research needs in different scenarios.

Description

Technical Field

[0001] The present invention relates to the field of image processing technology, and in particular to a Mars image augmentation method based on a stable diffusion model, a computing device, and a storage medium. Background Art

[0002] Due to the complex and ever-changing environment of Mars, often accompanied by dust and other weather conditions, obtaining high-resolution images of Mars is particularly difficult. Among them, Mars images are images of the Martian surface captured by a camera. Mars images include the topographic features of the Martian surface and are used for geological research on Mars.

[0003] However, combining deep learning models with geological research on Mars requires a large number of Mars images. Since the current data of Mars images is limited, it is necessary to use image augmentation to fully process a small number of Mars images to obtain a large number of Mars images.

[0004] Image augmentation is an important method and process for increasing the quantity and diversity of image data, helping to improve model generalization and performance. Related technologies typically use single-data deformation techniques to generate new images by applying different transformations to a single image. However, this approach suffers from insufficient diversity in the augmented images. Summary of the Invention

[0005] The present invention provides a Mars image augmentation method, computing device, and storage medium based on a stable diffusion model. The method can implement image augmentation processing of Mars images based on the diffusion model, and can quickly obtain large quantities of high-quality Mars images that meet image diversity requirements, thereby meeting users' needs for Mars geological research in different scenarios.

[0006] The technical solution of the present invention to solve the above technical problems is as follows:

[0007] In a first aspect, the present invention provides a method for augmenting a Mars image based on a stable diffusion model, the method comprising: obtaining a training sample set, the training sample set comprising multiple Mars images and explanatory information corresponding to each of the multiple Mars images, the explanatory information being used to describe the rock features and mineral features included in the Mars image; training the stable diffusion model according to the training sample set to obtain a trained stable diffusion model; performing conditional training on the trained stable diffusion model based on constraints to obtain a Control Net model, the Control Net model being used to control the image quality of the Mars image output by the trained stable diffusion model; obtaining an augmented image through the trained stable diffusion model and the Control Net model according to input feature information, the feature information comprising feature text and / or feature image.

[0008] In a possible implementation of the first aspect, obtaining an augmented image based on input feature information using a trained stable diffusion model and a Control Net model includes: obtaining a first latent vector based on the feature information using the trained stable diffusion model; obtaining a second latent vector based on the feature information using the Control Net model; summing the first latent vector and the second latent vector to obtain a target latent vector; and decoding the target latent vector to obtain the augmented image.

[0009] In a possible implementation of the first aspect, obtaining a first latent vector based on feature information using a trained stable diffusion model includes:

[0010] S301, encode the feature information to obtain a third latent vector;

[0011] S302, adding noise to the third latent vector to obtain a fourth latent vector;

[0012] S303: Perform noise prediction on the fourth latent vector to obtain first noise prediction information;

[0013] S304: Subtract the first noise prediction information from the third latent vector to obtain a fifth latent vector.

[0014] S305: Repeat steps S302-S304 according to the preset number of sampling times, and determine the fifth latent vector obtained last time as the first latent vector.

[0015] In a possible implementation of the first aspect, obtaining a second latent vector through a Control Net model according to feature information includes:

[0016] S401, encode the feature information to obtain a sixth latent vector;

[0017] S402, adding noise to the sixth latent vector to obtain a seventh latent vector;

[0018] S403: Perform noise prediction on the seventh latent vector according to the constraint condition to obtain second noise prediction information;

[0019] S404: Subtract the second noise prediction information from the sixth latent vector to obtain an eighth latent vector;

[0020] S405: Repeat steps S402-S404 according to the preset number of sampling times, and determine the eighth latent vector obtained last time as the second latent vector.

[0021] In a possible implementation of the first aspect, the trained stable diffusion model includes a text encoder, a variational self-decoding encoder, a UNet module, and a variational self-decoding decoder; the text encoder is used to encode the feature text to obtain a latent vector when the feature information includes feature text; the variational self-decoding encoder is used to encode the feature image to obtain a latent vector when the feature information includes a feature image; the UNet module is used to superimpose noise on the image; and to perform noise prediction on the image obtained by noise superposition to obtain noise prediction information; and also to subtract the noise prediction information from the image; the variational self-decoding decoder is used to decode the target latent vector to obtain an augmented image.

[0022] In a possible implementation of the first aspect, the Control Net model includes a text encoder, a variational self-decoding encoder, and a UNet module; the text encoder is used to encode the feature text to obtain a latent vector when the feature information includes feature text; the variational self-decoding encoder is used to encode the feature image to obtain a latent vector when the feature information includes a feature image; the UNet module is used to superimpose noise on the image; and to perform noise prediction on the image obtained by noise superposition according to constraints to obtain noise prediction information; and also to subtract the noise prediction information from the image.

[0023] In a possible implementation of the first aspect, before training the UNet module included in the trained stable diffusion model based on the constraints, the method further includes: receiving a sample image input by a user; and obtaining the constraints included in the sample image through Canny edge detection.

[0024] In a possible implementation of the first aspect, obtaining the constraints included in the sample image through Canny edge detection includes: smoothing the sample image through Gaussian filtering; determining the gradient value and gradient direction of each pixel included in the sample image; determining the pixel point whose gradient value is greater than a preset threshold and whose gradient value is the maximum value in the gradient direction as a candidate edge point; determining the candidate edge point whose gradient value is between a highest threshold and a lowest threshold as an edge point of the sample image; and determining the constraints included in the sample image based on the edge point of the sample image.

[0025] The beneficial effects of the present invention are as follows: the method provided by the present invention obtains a training sample set; trains a stable diffusion model based on the training sample set to obtain a trained stable diffusion model; conditionally trains the trained stable diffusion model based on constraints to obtain a Control Net model; and obtains an augmented image based on input feature information using the trained stable diffusion model and the Control Net model, wherein the feature information includes feature text and / or feature image. The method provided by the present invention implements image augmentation processing of Martian images based on the stable diffusion model, and can quickly obtain large quantities of high-quality Martian images that meet image diversity requirements, thereby meeting users' Martian geological research needs in different scenarios.

[0026] In a second aspect, a computing device is provided, comprising a memory and one or more processors; the memory is coupled to the processor; wherein computer program code is stored in the memory, and the computer program code comprises computer instructions, which, when executed by the processor, enable the computing device to execute a Mars image augmentation method based on a stable diffusion model as in any implementation of the first aspect.

[0027] In a third aspect, a computer-readable storage medium is provided, comprising computer instructions. When the computer instructions are executed on a computing device, the computing device executes a Mars image augmentation method based on a stable diffusion model as in any implementation of the first aspect.

[0028] In a fourth aspect, a computer program product is provided. When the computer program product is run on a computing device, the computing device is caused to execute the Mars image augmentation method based on the stable diffusion model as in any implementation of the first aspect.

[0029] It can be understood that the beneficial effects that can be achieved by the computing device described in the second aspect, the computer-readable storage medium described in the third aspect, and the computer program product described in the fourth aspect provided above can be referred to the beneficial effects in the first aspect and any possible design method thereof, and will not be repeated here. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 This is a schematic structural diagram of an augmentation device according to an embodiment of the present invention;

[0031] Figure 2 A schematic diagram of the structure of a computing device according to an embodiment of the present invention;

[0032] Figure 3 This is a flow chart of a Mars image augmentation method based on a stable diffusion model according to an embodiment of the present invention;

[0033] Figure 4A schematic diagram of a Mars image shown in an embodiment of the present invention;

[0034] Figure 5 This is a flow chart illustrating another Mars image augmentation method based on a stable diffusion model according to an embodiment of the present invention;

[0035] Figure 6 This is a flow chart illustrating another method for augmenting a Mars image based on a stable diffusion model according to an embodiment of the present invention;

[0036] Figure 7 This is a flow chart illustrating another method for augmenting a Mars image based on a stable diffusion model according to an embodiment of the present invention;

[0037] Figure 8 A schematic diagram of a Mars image augmentation method based on a stable diffusion model according to an embodiment of the present invention;

[0038] Figure 9a A schematic diagram of an augmented image according to an embodiment of the present invention;

[0039] Figure 9b This is a schematic diagram of another augmented image according to an embodiment of the present invention. DETAILED DESCRIPTION

[0040] The technical solutions in the embodiments of the present invention will be described below in conjunction with the accompanying drawings in the embodiments of the present invention. In the description of the embodiments of the present invention, unless otherwise specified, "multiple" refers to two or more than two. "At least one of the following" or similar expressions refers to any combination of these items, including any combination of single items or plural items. For example, at least one of a, b, or c can represent: a, b, c, ab, ac, bc, or abc, where a, b, and c can be single or multiple.

[0041] In addition, in order to facilitate the clear description of the technical solutions of the embodiments of the present invention, in the embodiments of the present invention, words such as "first" and "second" are used to distinguish between identical or similar items with substantially the same functions and effects. Those skilled in the art will understand that words such as "first" and "second" do not limit the quantity and execution order, and words such as "first" and "second" do not necessarily limit differences. At the same time, in the embodiments of the present invention, words such as "exemplary" or "for example" are used to indicate examples, illustrations or explanations. Any embodiment or design scheme described as "exemplary" or "for example" in the embodiments of the present invention should not be interpreted as being more preferred or advantageous than other embodiments or design schemes. Specifically, the use of words such as "exemplary" or "for example" is intended to present related concepts in a concrete way for easy understanding.

[0042] Due to the complex and ever-changing environment of Mars, often accompanied by dust and other weather conditions, obtaining high-resolution images of Mars is particularly difficult. Among them, Mars images are images of the Martian surface captured by a camera. Mars images include the topographic features of the Martian surface and are used for geological research on Mars.

[0043] However, combining deep learning models with geological research on Mars requires a large number of Mars images. Since the current data of Mars images is limited, it is necessary to use image augmentation to fully process a small number of Mars images to obtain a large number of Mars images.

[0044] Image augmentation is an important method and process for increasing the quantity and diversity of image data, helping to improve model generalization and performance. Related technologies typically use single-data deformation techniques to generate new images by applying different transformations to a single image. However, this approach suffers from insufficient diversity in the augmented images.

[0045] In view of this, an embodiment of the present invention provides a Mars image augmentation method based on a stable diffusion model, the method comprising: obtaining a training sample set, the training sample set comprising multiple Mars images and explanatory information corresponding to each of the multiple Mars images, the explanatory information being used to describe the rock features and mineral features included in the Mars image; training the stable diffusion model according to the training sample set to obtain a trained stable diffusion model; performing conditional training on the trained stable diffusion model based on constraints to obtain a Control Net model, the Control Net model being used to control the image quality of the Mars image output by the trained stable diffusion model; obtaining an augmented image through the trained stable diffusion model and the Control Net model according to input feature information, the feature information comprising feature text and / or feature image.

[0046] The method provided in an embodiment of the present invention obtains a training sample set; trains a stable diffusion model based on the training sample set to obtain a trained stable diffusion model; performs conditional training on the trained stable diffusion model based on constraints to obtain a Control Net model; and obtains an augmented image based on input feature information using the trained stable diffusion model and the Control Net model, wherein the feature information includes feature text and / or feature image. The method provided in an embodiment of the present invention implements image augmentation processing of Mars images based on the stable diffusion model, and can quickly obtain large quantities of high-quality Mars images that meet image diversity requirements, thereby meeting users' needs for Mars geological research in different scenarios.

[0047] In some embodiments, the Mars image augmentation method based on the stable diffusion model provided by the embodiments of the present invention can be executed by a Mars image augmentation device 100 based on the stable diffusion model (hereinafter referred to as augmentation device 100).

[0048] In one possible implementation, see Figure 1 , Figure 1 1 is a schematic structural diagram of a Mars image augmentation device according to an embodiment of the present invention. The Mars image augmentation device 100 includes a trained stable diffusion model 110 and a ControlNet model 120.

[0049] Among them, the trained stable diffusion model includes a text encoder, a variational self-decoding encoder, a UNet module and a variational self-decoding decoder; the text encoder is used to encode the feature text to obtain a latent vector when the feature information includes feature text, the variational self-decoding encoder is used to encode the feature image to obtain a latent vector when the feature information includes feature image, the UNet module is used to superimpose noise on the image; and to perform noise prediction on the image obtained by noise superposition to obtain noise prediction information; and also to subtract the noise prediction information from the image; the variational self-decoding decoder is used to decode the target latent vector to obtain an augmented image.

[0050] The Control Net model includes a text encoder, a variational self-decoding encoder, and a UNet module; the text encoder is used to encode the feature text to obtain a latent vector when the feature information includes feature text; the variational self-decoding encoder is used to encode the feature image to obtain a latent vector when the feature information includes feature image; the UNet module is used to superimpose noise on the image; and to perform noise prediction on the image obtained by noise superposition according to constraint conditions to obtain noise prediction information; and also to subtract the noise prediction information from the image.

[0051] This can also be understood as follows: the text encoder and variational self-decoding encoder convert the input noise vector into a compressed latent representation vector through a series of convolutional layers and nonlinear activation functions. The variational self-decoding decoder restores the latent representation vector to the image space dimension and combines it with the feature map of the corresponding encoder layer through skip connections. In the reverse process of the diffusion algorithm, the difference between the output image and the original image is set as the loss term. Using the U-Net method, the noise is learned iteratively to minimize the loss.

[0052] In one example, the UNet module includes 8 sampling layers and 17 main parameter units, each of which includes 4 ResNet layers and 2 Vi T modules. Each Vi T module contains several cross-attention modules or self-attention modules, and the diffusion time adopts position encoding.

[0053] As an example, the augmentation device 100 can be any computing device 200 with data processing capabilities, such as a general-purpose computer, a personal computer, a laptop computer, a switch, or a tablet computer. The specific implementation method of the augmentation device 100 is not limited here.

[0054] Figure 2 FIG2 shows a hardware structure diagram of a computing device 200 provided by an embodiment of the present invention. The computing device 200 includes a processor 210 , a memory 220 , and a communication interface 230 .

[0055] The processor 210 may include one or more processing cores. The processor 210 connects various components within the computing device 200 using various interfaces and lines. It executes various functions of the computing device 200 and processes data by running or executing instructions, programs, code sets, or instruction sets stored in the memory 220, and by calling data stored in the memory 220. Optionally, the processor 210 may be implemented in the form of at least one of a central processing unit (CPU), a graphics processing unit (GPU), a digital signal processing unit (DSP), a field-programmable gate array (FPGA), and a programmable logic array (PLA).

[0056] The memory 220 may include a random access memory (RAM) or a read-only memory (ROM). Optionally, the memory 220 includes a non-transitory computer-readable storage medium. The memory 220 may be used to store instructions, programs, codes, code sets, or instruction sets. The memory 220 may include a program storage area. The program storage area may store instructions for implementing an operating system, instructions for implementing at least one function (such as a data acquisition function, a model training function, an image generation function, etc.), instructions for implementing the above-mentioned various method embodiments, etc.

[0057] The communication interface 230 is used to communicate with other devices, equipment or communication networks, such as data storage devices, image processing equipment or Ethernet, radio access network (RAN), wireless local area network (WLAN), etc.

[0058] In physical implementation, the aforementioned components (e.g., processor 210, memory 220, and communication interface 230) may be components within the same device (e.g., a laptop). Alternatively, at least two of the components may be provided within the same device, i.e., as different components within a single device, similar to the deployment of devices or components in a distributed system.

[0059] It should be understood that the structure illustrated in this embodiment does not constitute a specific limitation on the computing device 200. In other embodiments of the present invention, the computing device 200 may include more or fewer components than shown, or may combine or separate certain components, or arrange the components differently. The components shown in the figure may be implemented in hardware, software, or a combination of software and hardware.

[0060] The following describes a Mars image augmentation method based on a stable diffusion model provided by an embodiment of the present invention in conjunction with the accompanying drawings.

[0061] Figure 3 The flowchart of a Mars image augmentation method based on a stable diffusion model provided by an embodiment of the present invention. Optionally, the method can be performed by Figure 1 The augmentation device 100 of the hardware structure shown is executed by Figure 2 The method may include the following steps:

[0062] S1. Obtain a training sample set, where the training sample set includes multiple Mars images and description information corresponding to each of the multiple Mars images, where the description information is used to describe rock features and mineral features included in the Mars image.

[0063] The training sample set includes multiple Mars images, each of which has the same resolution.

[0064] It should be noted that the captions for each Mars image serve as triggers for the stable diffusion model's training. These are labels corresponding to each Mars image, representing the image's elemental features. For example, they include the rock and mineral types and colors, the positional relationships between different rock and mineral types, and the topographical information at different locations within each Mars image.

[0065] For example, see Figure 4 , Figure 4 Schematic diagram of a Mars image A shown in the present invention, wherein the Mars image A includes basalt, and the description information of the Mars image A is “rocky area, the foreground includes rocks and gravel, and the background includes mountains and deserts”.

[0066] S2. Train the stable diffusion model according to the training sample set to obtain a trained stable diffusion model.

[0067] Specifically, the stable diffusion model is fine-tuned based on the training sample set and then trained to obtain a trained stable diffusion model. This method of fine-tuning the stable diffusion model and then training it is called the LORA method, and the trained stable diffusion model can also be called a LORA model.

[0068] Among them, the LORA model is based on the stable diffusion model and fine-tunes the model parameters of the stable diffusion model so that it can adapt to various application scenarios for rapid training, thereby improving the convergence speed and portability of the model.

[0069] Exemplarily, the augmentation device 100 inputs the training sample set and model parameters into the stable diffusion model. The stable diffusion model is trained based on the training sample set and the model parameters to obtain a trained stable diffusion model. The model parameters include the image resolution and number of Mars images included in the training sample set, the step size and number of iterations per iteration, and the optimizer type of the stable diffusion model.

[0070] It can also be understood that the LORA model has the same network structure as the stable diffusion model, and the augmentation device 100 adjusts the cross-attention layer of the UNet module included in the stable diffusion model according to the model parameters and the training sample set. During the training process of the stable diffusion model, the other layers of the stable diffusion model are kept unchanged, and the cross-attention layer of the stable diffusion model is fine-tuned based on the model parameters. Then, the cross-attention layer of the fine-tuned LoRA model is combined with the other layers of the stable diffusion model to obtain a trained stable diffusion model, that is, a LORA model. The trained stable diffusion model is used to obtain the corresponding augmented image based on the input feature information.

[0071] In one example, the stable diffusion model is in the latent space. The forward process of the diffusion algorithm gradually adds noise to a given image z0 in t steps until a pure noise image z is generated. t , combined with text prompt c t and task control condition c f , the model learns a network ∈ θ The objective function of learning noise and stable diffusion model can be expressed as:

[0072]

[0073] Among them, ∈~N(0,1) indicates that the noise distribution belongs to Gaussian distribution.

[0074] S3. Based on the constraint conditions, conditional training is performed on the trained stable diffusion model to obtain a Control Net model. The Control Net model is used to control the image quality of the Mars image output by the trained stable diffusion model.

[0075] Specifically, the ControlNet model is also based on a stable diffusion model and has a network structure similar to that of the stable diffusion model. The ControlNet model trains the encoder and decoder included in the stable diffusion model based on constraints to optimize the encoding and decoding results of the encoder and decoder included in the stable diffusion model. It can also be understood that ControlNet is a neural network structure that controls the diffusion model by controlling input constraints. Constraints are input during the generation process to control the properties of the augmented image to improve the image quality of the augmented image, thereby achieving more fine-grained control of the image quality of the augmented image provided by the embodiments of the present invention, thereby making the style and quality of the augmented image closer to the Mars image included in the training sample set.

[0076] In a possible implementation, before the above S3, the method provided by the embodiment of the present invention further includes the following steps: receiving a sample image input by a user; and obtaining constraints included in the sample image through Canny edge detection.

[0077] Optionally, constraints on the sample image obtained through Canny edge detection include:

[0078] The sample image is smoothed by Gaussian filtering; the gradient value and gradient direction of each pixel point included in the sample image are determined; the pixel point whose gradient value is greater than a preset threshold and whose gradient value is the maximum in the gradient direction is determined as a candidate edge point; the candidate edge point whose gradient value is between the highest threshold and the lowest threshold is determined as the edge point of the sample image; and the constraint conditions included in the sample image are determined based on the edge points of the sample image.

[0079] It should be understood that the method provided in the embodiment of the present invention can also input constraints in other ways. The above method of obtaining constraints by Canny edge detection is only an example, and the embodiment of the present invention does not impose any special restrictions on this.

[0080] To facilitate understanding, the specific training process of the Control Net model is illustrated below. The augmentation device 100 copies the trained stable diffusion model 110, which includes encoders A1, A2, and A3. The augmentation device 100 conditionally trains encoders A1, A2, and A3 according to constraints to obtain encoders B1, B2, and B3. Finally, encoders B1, B2, and B3 are combined with the portion of the trained stable diffusion model 110 excluding encoders A1, A2, and A3 to obtain the Control Net model.

[0081] S4. Obtain an augmented image through the trained stable diffusion model and Control Net model according to the input feature information, where the feature information includes feature text and / or feature image.

[0082] In one possible implementation, see Figure 5 The above S4 specifically includes the following steps:

[0083] S51. Obtain a first latent vector through a trained stable diffusion model based on the feature information.

[0084] Optional, see Figure 6 The above S51 specifically includes the following steps:

[0085] S301, encode the feature information to obtain a third latent vector;

[0086] S302, adding noise to the third latent vector to obtain a fourth latent vector;

[0087] S303: Perform noise prediction on the fourth latent vector to obtain first noise prediction information;

[0088] S304: Subtract the first noise prediction information from the third latent vector to obtain a fifth latent vector.

[0089] S305: Repeat steps S302-S304 according to the preset number of sampling times, and determine the fifth latent vector obtained last time as the first latent vector.

[0090] S52. Obtain a second latent vector through the Control Net model according to the feature information.

[0091] Optional, see Figure 7 The above S52 specifically includes the following steps:

[0092] S401, encode the feature information to obtain a sixth latent vector;

[0093] S402, adding noise to the sixth latent vector to obtain a seventh latent vector;

[0094] S403: Perform noise prediction on the seventh latent vector according to the constraint condition to obtain second noise prediction information;

[0095] S404: Subtract the second noise prediction information from the sixth latent vector to obtain an eighth latent vector;

[0096] S405. Repeat steps S402-S404 according to the preset number of sampling times, and determine the eighth latent vector obtained last time as the second latent vector.

[0097] S53. Add the first latent vector and the second latent vector to obtain a target latent vector.

[0098] S54. Decode the target latent vector to obtain an augmented image.

[0099] In a possible implementation, the feature information further includes image resolution, number of images, number of generation steps, positive prompt words, and negative prompt words of the augmented image.

[0100] It should be understood that the above description is merely exemplary, and the embodiments of the present invention do not impose any particular limitation on the specific content of the feature information.

[0101] In order to facilitate understanding of this solution, the following is a specific process of the augmentation device 100 obtaining the augmented image through the stable diffusion model and the Control Net model trained according to the input feature information. Figure 8 For example, the augmentation device 100 receives input feature information and inputs the feature information into the trained stable diffusion model 110 and the Control Net model 120. The trained stable diffusion model 110 obtains a first latent vector based on the input feature information, and the Control Net model 120 obtains a second latent vector based on the input feature information. The augmentation device 100 then sums the first latent vector and the second latent vector to obtain a target latent vector. Finally, the target latent vector is decoded by a decoder to obtain an augmented image.

[0102] For example, see Figure 9a and Figure 9b , Figure 9a is the augmented image output by the augmentation device 100 through the trained stable diffusion model based on the input feature information, Figure 9a The augmented image is output by the augmentation device 100 through the stable diffusion model and the Control Net model trained according to the input feature information. Figure 9a and Figure 9bIt can be seen that the ControlNet model is able to control the quality of the augmented images and achieve finer-grained control over the image quality of the augmented images, thereby making the style and quality of the augmented images closer to the Mars images included in the training sample set.

[0103] As can be seen from S1-S4 above, the method provided by the embodiment of the present invention implements image augmentation processing for Mars images based on a stable diffusion model, enabling rapid production of large quantities of high-quality Mars images that meet image diversity requirements, thereby satisfying users' needs for Martian geological research in various scenarios. Specifically, the method provided by the embodiment of the present invention implements Martian image augmentation using a trained stable diffusion model and controls the quality of the augmented images based on a Control Net model, thereby ensuring that the style and quality of the augmented images more closely resemble the Mars images included in the training sample set.

[0104] The above mainly introduces the solution of the embodiment of the present invention from the perspective of method. It can be understood that in order to realize the above functions, the augmentation device 100 includes at least one of the hardware structure and software modules corresponding to the execution of each function. It should be easy for those skilled in the art to realize that, in combination with the units and algorithm steps of each example described in the embodiment disclosed herein, the embodiment of the present invention can be implemented in the form of hardware or a combination of hardware and computer software. Whether a function is executed in the form of hardware or computer software driving hardware depends on the specific application and design constraints of the technical solution. Professional and technical personnel can use different methods to implement the described functions for each specific application, but such implementation should not be considered to exceed the scope of the embodiment of the present invention.

[0105] An embodiment of the present invention further provides a computer-readable storage medium storing at least one computer instruction, which is loaded and executed by a processor to implement the methods of each of the above embodiments. For explanations of the relevant contents and descriptions of the beneficial effects of any of the above-mentioned computer-readable storage media, reference can be made to the corresponding embodiments described above and will not be repeated here.

[0106] The embodiment of the present invention further provides a chip. The chip integrates a control circuit and one or more ports for implementing the functions of the above-mentioned augmentation device 100. Optionally, the functions supported by the chip can be referred to above and will not be repeated here.

[0107] Those skilled in the art will appreciate that all or part of the steps of the above-mentioned embodiments can be implemented by a program to instruct the relevant hardware to complete the program, which can be stored in a computer-readable storage medium. The above-mentioned storage medium can be a read-only memory, a random access memory, etc. The above-mentioned processing unit or processor can be a central processing unit, a general-purpose processor, a specific circuit structure (appliance specific integrated circuit, ASIC), a microprocessor (digital integrated processor, DSP), a field programmable gate array (field programmable gate array, FPGA) or other programmable logic devices, transistor logic devices, hardware components or any combination thereof.

[0108] The embodiment of the present invention further provides a computer program product comprising instructions, which, when executed on a computer, causes the computer to perform any of the methods in the above embodiments. The computer program product comprises one or more computer instructions. When the computer program instructions are loaded and executed on a computer, the process or function according to the embodiment of the present invention is generated in whole or in part. The computer may be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions may be stored in a computer-readable storage medium or transmitted from one computer-readable storage medium to another computer-readable storage medium. For example, the computer instructions may be transmitted from one website, computer, server, or data center to another website, computer, server, or data center via a wired (e.g., coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) method. The computer-readable storage medium may be any available medium that a computer can access or a data storage device such as a server or data center that includes one or more available media integrated therein. The available medium may be a magnetic medium (e.g., a floppy disk, a hard disk, a magnetic tape), an optical medium (e.g., a DVD), or a semiconductor medium (e.g., an SSD).

[0109] It should be noted that the above-mentioned devices for storing computer instructions or computer programs provided in the embodiments of the present invention, such as but not limited to the above-mentioned memories, computer-readable storage media and communication chips, etc., are all non-transitory. Those skilled in the art should be aware that in one or more of the above examples, the functions described in the embodiments of the present invention can be implemented using hardware, software, firmware or any combination thereof. When implemented using software, these functions can be stored in a computer-readable storage medium or transmitted as one or more instructions or codes on a computer-readable storage medium. Computer-readable storage media include computer storage media and communication media, wherein the communication medium includes any medium that facilitates the transmission of computer programs from one place to another. The storage medium can be any available medium that can be accessed by a general-purpose or special-purpose computer.

[0110] Although the embodiments of the present invention have been shown and described above, it will be understood that the above embodiments are illustrative and are not to be construed as limitations on the present invention. A person skilled in the art may change, modify, replace and modify the above embodiments within the scope of the present invention.

Claims

1. A Mars image augmentation method based on a stable diffusion model, characterized in that: The method comprises: Obtaining a training sample set, the training sample set comprising a plurality of Mars images and explanatory information corresponding to each of the plurality of Mars images, the explanatory information being used to describe rock features and mineral features included in the Mars image; wherein the explanatory information is a geological semantic annotation of each Mars image, including: rock and mineral type, color characteristics, spatial position relationship, and landform information; the explanatory information is a trigger word for a stable diffusion model; Training the stable diffusion model according to the training sample set to obtain a trained stable diffusion model; Based on the constraint conditions, conditional training is performed on the trained stable diffusion model to obtain a ControlNet model, where the ControlNet model is used to control the image quality of the Mars image output by the trained stable diffusion model; Obtaining an augmented image through the trained stable diffusion model and the ControlNet model according to input feature information, wherein the feature information includes feature text and / or feature image; Obtaining an augmented image through the trained stable diffusion model and the ControlNet model according to the input feature information includes: Obtaining a first latent vector according to the feature information through the trained stable diffusion model; Obtaining a second latent vector through the ControlNet model according to the feature information; Adding the first latent vector and the second latent vector to obtain a target latent vector; Decoding the target latent vector to obtain the augmented image; The trained stable diffusion model includes a text encoder, a variational self-decoding encoder, a UNet module, and a variational self-decoding decoder; the text encoder is used to encode the feature text to obtain a latent vector when the feature information includes feature text; the variational self-decoding encoder is used to encode the feature image to obtain a latent vector when the feature information includes a feature image; the UNet module is used to perform noise superposition on the image; and to perform noise prediction on the image obtained by noise superposition to obtain noise prediction information; and further to subtract the noise prediction information from the image; the variational self-decoding decoder is used to decode the target latent vector to obtain the augmented image; The ControlNet model includes a text encoder, a variational self-decoding encoder and a UNet module; the text encoder is used to encode the feature text to obtain a latent vector when the feature information includes feature text; the variational self-decoding encoder is used to encode the feature image to obtain a latent vector when the feature information includes a feature image; the UNet module is used to superimpose noise on the image; and to perform noise prediction on the image obtained by noise superposition according to the constraint conditions to obtain noise prediction information; and also to subtract the noise prediction information from the image.

2. The method according to claim 1, characterized in that Obtaining a first latent vector according to the feature information through the trained stable diffusion model includes: S301, encoding the feature information to obtain a third latent vector; S302: Add noise to the third latent vector to obtain a fourth latent vector; S303: Perform noise prediction on the fourth latent vector to obtain first noise prediction information; S304: Subtract the first noise prediction information from the third latent vector to obtain a fifth latent vector; S305: Repeat steps S302-S304 according to a preset number of sampling times, and determine the fifth latent vector obtained last time as the first latent vector.

3. The method according to claim 2, characterized in that Obtaining a second latent vector through the ControlNet model according to the feature information includes: S401, encoding the feature information to obtain a sixth latent vector; S402: Add noise to the sixth latent vector to obtain a seventh latent vector; S403: Perform noise prediction on the seventh latent vector according to the constraint condition to obtain second noise prediction information; S404: Subtract the second noise prediction information from the sixth latent vector to obtain an eighth latent vector; S405: Repeat steps S402-S404 according to a preset number of sampling times, and determine the eighth latent vector obtained last time as the second latent vector.

4. The method according to claim 1, wherein Before training the UNet module included in the trained stable diffusion model based on the constraint conditions, the method further includes: Receive a sample image input by the user; The constraints included in the sample image are obtained through Canny edge detection.

5. The method according to claim 4, characterized in that The constraints involved in acquiring the sample image through Canny edge detection include: Performing smoothing processing on the sample image by Gaussian filtering; Determine the gradient value and gradient direction of each pixel point included in the sample image; The pixel whose gradient value is greater than the preset threshold and whose gradient value is the maximum value in the gradient direction is determined as a candidate edge point; Determine a candidate edge point whose gradient value is between a maximum threshold and a minimum threshold as an edge point of the sample image; The constraint conditions included in the sample image are determined according to the edge points of the sample image.

6. A computing device, characterized in that include: processor; a memory for storing instructions executable by the processor; The processor is configured to execute the instructions to implement the Mars image augmentation method based on the stable diffusion model as described in any one of claims 1 to 5.

7. A computer-readable medium having computer program instructions stored thereon, wherein the computer program instructions can be executed by a processor to implement the Mars image augmentation method based on the stable diffusion model as described in any one of claims 1 to 5.

Images