Method and apparatus for training an image generation model
By employing a divide-and-conquer strategy in the image generation model, and utilizing generative adversarial networks and fully convolutional networks to process background and text style transfer, combined with a pre-trained text recognition model, the problem of insufficient generalization ability of optical character recognition models for rare characters is solved, achieving efficient data generation and improved recognition performance.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- HARBIN INST OF TECH SHENZHEN GRADUATE SCHOOL
- Filing Date
- 2022-11-01
- Publication Date
- 2026-06-02
AI Technical Summary
Existing optical character recognition models lack the ability to generalize to rare characters in the training dataset, resulting in poor recognition performance and high dataset collection costs.
An image generation model employing a divide-and-conquer strategy processes background and text styles through a background erasure module and a text conversion module, respectively. It utilizes generative adversarial networks and fully convolutional networks for image generation and text conversion, and combines a pre-trained text recognition model for supervision to generate the target image.
It improves the generalization ability of optical character recognition models to rare characters, reduces the cost of dataset collection, and generates images that maintain data consistency and have good text readability.
Smart Images

Figure CN116091630B_ABST
Abstract
Description
Technical Field
[0001] The embodiments of this disclosure relate to the field of computer technology, and more specifically, to a method and apparatus for training an image generation model. Background Technology
[0002] Text is a crucial visual cue, providing rich and accurate high-level semantic information. As a carrier of semantic information, text can supplement and enhance the information conveyed by an image itself. In the digital age, optical character recognition (OCR) has become an important means of understanding image or video content in order to convert text in images into computer-editable text. However, in the era of deep learning, the effectiveness of OCR models largely depends on the training data in the training dataset. On the one hand, collecting training data requires significant human and material resources. On the other hand, in some languages with large character sets (such as Chinese), the data is often long-tailed. That is, real-world datasets often contain rare characters, and the presence of these characters greatly affects the generalization of OCR models. Therefore, data augmentation techniques are often used in scene text recognition tasks to obtain the desired training dataset.
[0003] Image generation methods based on text style transfer can modify the text content in a style image to the target text content without altering the overall style of the style image, such as background texture, font color, and font shape. Therefore, this image generation method can be used as a data augmentation method for optical character recognition. This allows for the generation of multiple new images containing various characters (including rare characters) from real images, while maintaining good data consistency between the real images and the new images. Summary of the Invention
[0004] The embodiments described herein provide a method, apparatus, and computer-readable storage medium storing a computer program for training an image generation model.
[0005] According to a first aspect of this disclosure, a method for training an image generation model is provided. The image generation model includes a background erasure module, a text conversion module, and a fusion module. The method includes: acquiring a training image set, wherein the training image set includes: a style image containing first text and a content image containing second text. The first text has a first text style attribute, and the second text has a second text style attribute; obtaining a background image and a first text mask image corresponding to the background image from the style image through the background erasure module, wherein the first text mask image is used to calculate a loss function of the background erasure module; converting the second text style attribute of the second text in the content image to the first text style attribute through the text conversion module to obtain a converted content image and a second text mask image corresponding to the converted content image, wherein the second text mask image is used to calculate a loss function of the text conversion module; fusing the converted content image with the background image through the fusion module to obtain a target image; and minimizing the loss function of the background erasure module, the loss function of the text conversion module, and the loss function of the fusion module to update the parameters of the image generation model.
[0006] In some embodiments of this disclosure, the background erasure module includes a first generative adversarial network (GAN). The first generator in the first GAN employs a UNet-like encoder-decoder structure. Obtaining a background image and a corresponding first text mask image from a style image via the background erasure module includes: encoding the style image using a first encoder in the first generator to obtain a latent space representation of the style image; reconstructing the latent space representation into a background image using a first decoder symmetrical to the first encoder in the first generator; and reconstructing the latent space representation into a first text mask image using a second decoder symmetrical to the first encoder in the first generator.
[0007] In some embodiments of this disclosure, the training image set further includes: a real background image corresponding to the style image, and a real text mask image corresponding to the style image. The loss function of the background erasure module is equal to the sum of the first loss function, the second loss function, and the third loss function. The first loss function is expressed as:
[0008]
[0009] Among them, D B I represents the first discriminator in the first generative adversarial network. s Represents style image, O b T represents the background image. b This represents the real background image corresponding to the style image.
[0010] The second loss function is expressed as:
[0011]
[0012] Among them, O ms T represents the first text mask image. ms This represents the real text mask image corresponding to the style image, where N represents the size of the first text mask image.
[0013] The third loss function is expressed as:
[0014] L L1_1 =||T b -O b ||1,
[0015] Among them, O b T represents the background image. b This represents the real background image corresponding to the style image.
[0016] In some embodiments of this disclosure, the method further includes: supervising the text readability of the converted content image using a pre-trained text recognition model.
[0017] In some embodiments of this disclosure, the text conversion module includes an encoder-decoder structure based on a fully convolutional network. Converting a second text style attribute of second text in a content image to a first text style attribute using the text conversion module to obtain a converted content image and a corresponding second text mask image includes: encoding the style image using a second encoder in the encoder-decoder structure to obtain style image features; encoding the content image using a third encoder in the encoder-decoder structure to obtain content image features; concatenating the style image features and content image features into a fused feature block along the channel dimension; reconstructing the converted content image using the fused feature block using the third decoder in the encoder-decoder structure; and reconstructing the second text mask image using the fused feature block using the fourth decoder in the encoder-decoder structure.
[0018] In some embodiments of this disclosure, the training image set further includes: a real text mask image corresponding to the converted content image, and a real converted content image corresponding to the converted content image. The loss function of the text conversion module is equal to the sum of the fourth loss function, the fifth loss function, and the sixth loss function. The fourth loss function is expressed as:
[0019]
[0020] Where M represents the number of characters in the dictionary used by the text recognition model, k represents the number of characters in the second text in the converted content image, and p icy represents the predicted probability that the i-th character in the second text of the transformed content image is the c-th character in the dictionary. ic This indicates whether the i-th character in the second text of the converted content image is the c-th character in the dictionary. If y ic A value of 1 indicates that the i-th character in the second text of the converted content image is the c-th character in the dictionary.
[0021] The fifth loss function is expressed as:
[0022]
[0023] Among them, O mt T represents the second text mask image. mt This represents the real text mask image corresponding to the converted content image, where n represents the size of the second text mask image.
[0024] The sixth loss function is expressed as:
[0025] L L1_2 =||T t -O t ||1,
[0026] Among them, O t T represents the converted content image. t This represents the actual converted content image that corresponds to the converted content image.
[0027] In some embodiments of this disclosure, the fusion module includes a second generative adversarial network (GAN). The second generator in the second GAN includes a first convolutional layer and a second convolutional layer. The fusion module fuses the transformed content image with the background image to obtain a target image, comprising: stitching the background image and the transformed content image into a stitched image along the channel dimension; reducing the dimensionality of the stitched image and obtaining fused image features through the first convolutional layer; and interacting with feature information on the fused image features through the second convolutional layer to generate the target image.
[0028] In some embodiments of this disclosure, the method further includes supervising the text readability of the target image using a pre-trained text recognition model.
[0029] In some embodiments of this disclosure, the training image set further includes: a real target image corresponding to the target image. The method further includes: using an image feature extraction model to extract features from the target image and the real target image corresponding to the target image. The loss function of the fusion module is equal to the sum of the seventh, eighth, ninth, tenth, and eleventh loss functions. The seventh loss function is expressed as:
[0030]
[0031] Where M represents the number of characters in the dictionary used by the text recognition model, k represents the number of characters in the second text in the target image, and P ic Y represents the predicted probability that the i-th character in the second text of the target image is the c-th character in the dictionary. ic Indicate whether the i-th character in the second text of the target image is the c-th character in the dictionary. If Y ic A value of 1 indicates that the i-th character in the second text of the target image is the c-th character in the dictionary.
[0032] The eighth loss function is expressed as:
[0033]
[0034] Where, φ i (·) represents the i-th feature map output by the image feature extraction model, O f T represents the target image. f represents the real target image corresponding to the target image, and h represents the number of feature maps output by the image feature extraction model.
[0035] The ninth loss function is expressed as:
[0036]
[0037] Among them, G i (·) represents the Gram matrix of the i-th feature map output by the image feature extraction model, O f T represents the target image. f M represents the real target image corresponding to the target image, h represents the number of feature maps output by the image feature extraction model, and M represents the number of feature maps. i This represents the size of the i-th feature map.
[0038] The tenth loss function is expressed as:
[0039]
[0040] Among them, D F I represents the second discriminator in the second generative adversarial network. t Represents the content image, O f T represents the target image. f This represents the real target image corresponding to the target image.
[0041] The eleventh loss function is expressed as:
[0042] L L1_3 =||Tf -O f ||1,
[0043] Among them, O f T represents the target image. f This represents the real target image corresponding to the target image.
[0044] According to a second aspect of this disclosure, an apparatus for training an image generation model is provided. The apparatus includes at least one processor and at least one memory storing a computer program. When the computer program is executed by the at least one processor, the apparatus causes to: acquire a training image set, wherein the training image set includes: a style image containing first text and a content image containing second text, the first text having a first text style attribute and the second text having a second text style attribute; obtain a background image and a first text mask image corresponding to the background image from the style image via a background erasure module, wherein the first text mask image is used to calculate a loss function of the background erasure module; convert the second text style attribute of the second text in the content image to the first text style attribute via a text conversion module to obtain a converted content image and a second text mask image corresponding to the converted content image, wherein the second text mask image is used to calculate a loss function of the text conversion module; fuse the converted content image with the background image via a fusion module to obtain a target image; and minimize the loss function of the background erasure module, the loss function of the text conversion module, and the loss function of the fusion module to update the parameters of the image generation model.
[0045] In some embodiments of this disclosure, the background erasure module includes a first generative adversarial network (GAN). The first generator in the first GAN employs a UNet-shaped encoder-decoder structure. When executed by at least one processor, the computer program causes the apparatus to obtain a background image and a corresponding first text mask image from a style image via the background erasure module by: encoding the style image using a first encoder in the first generator to obtain a latent space representation of the style image; reconstructing the latent space representation into a background image using a first decoder symmetrical to the first encoder in the first generator; and reconstructing the latent space representation into a first text mask image using a second decoder symmetrical to the first encoder in the first generator.
[0046] In some embodiments of this disclosure, when executed by at least one processor, the computer program causes the apparatus to further: supervise the text readability of the converted content image by means of a pre-trained text recognition model.
[0047] In some embodiments of this disclosure, the text conversion module includes an encoder-decoder structure based on a fully convolutional network. When executed by at least one processor, a computer program causes the apparatus to convert a second text style attribute of second text in a content image into a first text style attribute via the text conversion module to obtain a converted content image and a second text mask image corresponding to the converted content image by: encoding the style image using a second encoder in the encoder-decoder structure to obtain style image features; encoding the content image using a third encoder in the encoder-decoder structure to obtain content image features; concatenating the style image features and content image features into a fused feature block in the channel dimension; reconstructing the converted content image using the fused feature block using the third decoder in the encoder-decoder structure; and reconstructing the second text mask image using the fused feature block using a fourth decoder in the encoder-decoder structure.
[0048] In some embodiments of this disclosure, the fusion module includes a second generative adversarial network (GAN). The second generator in the second GAN includes a first convolutional layer and a second convolutional layer. When executed by at least one processor, a computer program causes the apparatus to fuse a transformed content image with a background image through the fusion module to obtain a target image by: stitching the background image and the transformed content image into a stitched image in the channel dimension; reducing the dimensionality of the stitched image through the first convolutional layer and obtaining fused image features; and interacting with feature information on the fused image features through the second convolutional layer to generate the target image.
[0049] In some embodiments of this disclosure, when executed by at least one processor, the computer program causes the device to further: supervise the text readability of the target image by means of a pre-trained text recognition model.
[0050] In some embodiments of this disclosure, the training image set further includes: a real target image corresponding to the target image. When executed by at least one processor, the computer program causes the apparatus to further: use an image feature extraction model to extract features from the target image and the real target image corresponding to the target image.
[0051] According to a third aspect of this disclosure, a computer-readable storage medium is provided storing a computer program, wherein the computer program, when executed by a processor, implements the steps of the method described according to a first aspect of this disclosure. Attached Figure Description
[0052] To more clearly illustrate the technical solutions of the embodiments of this disclosure, the accompanying drawings of the embodiments will be briefly described below. It should be understood that the drawings described below only relate to some embodiments of this disclosure and are not intended to limit this disclosure, wherein:
[0053] Figure 1 This is an exemplary flowchart of a method for training an image generation model according to embodiments of the present disclosure;
[0054] Figure 2 This is a schematic block diagram of an image generation model according to an embodiment of the present disclosure;
[0055] Figure 3 It is a schematic diagram of three sets of style images, content images, and the target image derived from them; and
[0056] Figure 4 This is a schematic block diagram of an apparatus for training an image generation model according to embodiments of the present disclosure.
[0057] In the accompanying diagram, markers with the same last two digits correspond to the same elements. It should be noted that the elements in the diagram are schematic and not drawn to scale. Detailed Implementation
[0058] To make the objectives, technical solutions, and advantages of the embodiments of this disclosure clearer, the technical solutions of the embodiments of this disclosure will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this disclosure. All other embodiments obtained by those skilled in the art based on the described embodiments of this disclosure without creative effort are also within the scope of protection of this disclosure.
[0059] Unless otherwise defined, all terms used herein (including technical and scientific terms) shall have the same meaning as commonly understood by one of ordinary skill in the art to which this subject matter pertains. It will be further understood that terms such as those defined in commonly used dictionaries shall be interpreted as having meanings consistent with their meanings in the context of the specification and in the relevant art, and shall not be interpreted in an idealized or overly formal form unless otherwise explicitly defined herein. Furthermore, terms such as “first” and “second” are used only to distinguish one component (or part of a component) from another component (or another part of a component).
[0060] The main challenges of text style transfer-based image generation tasks lie in text style transfer and background texture reconstruction. On the one hand, text style is a complex mixture, including elements such as font, color, size, and position, making text style extraction and transfer very challenging. On the other hand, styled images often contain common quality issues such as shadows, blur, and wrinkles, making background erasure and reconstruction even more difficult.
[0061] This disclosure proposes an image generation model. The model employs a divide-and-conquer strategy, utilizing two separate modules (a background erasure module and a text conversion module) to implement the background erasure and text conversion functions in a text style transfer-based image generation task, respectively. The background erasure module erases text regions from the style image and restores the background information of those regions, generating a clean background image. The text conversion module extracts the text style from the style image and applies this style to the text content in the content image, generating a converted content image. The image generation model then uses a fusion module to fuse the background image output by the background erasure module and the converted content image output by the text conversion module to achieve the target image. Thus, the target image will present the background from the style image and the text content of the converted style from the content image.
[0062] Embodiments of this disclosure also propose a method for training an image generation model. Figure 1 An exemplary flowchart of a method 100 for training an image generation model according to an embodiment of the present disclosure is shown. Figure 2 A schematic block diagram of an image generation model 200 according to an embodiment of the present disclosure is shown. The image generation model 200 may include: a background erasure module 210, a text conversion module 220, and a fusion module 230. The following is in conjunction with... Figure 1 and Figure 2 This describes the process of training an image generation model.
[0063] exist Figure 1 At box S102, the training image set is obtained. The training image set may include: a style image containing the first text and a content image containing the second text. The first text has a first text style attribute. The second text has a second text style attribute. Figure 2 In the example, style image I s The text content is "durian pizza", and the style image is I. s The first text in the document is in Founder ShuTi font, with a font size of 22pt. (Content Image I) t The text content is "Checkout List", and the content image is I. t The second text in the document is in SimSun font and has a font size of 26pt.
[0064] In some embodiments of this disclosure, in order to calculate the loss function of the image generation model, the training image set may further include: a real background image corresponding to the style image, a real text mask image corresponding to the style image, a real text mask image corresponding to the converted content image, a real converted content image corresponding to the converted content image, and a real target image corresponding to the target image.
[0065] To better generate the target image, embodiments of this disclosure rely on pairs of images before and after text style conversion. In real-life scenarios, obtaining paired image data is difficult, and it is usually obtained through image editing software (e.g., Photoshop). To reduce the difficulty of obtaining training data, in some embodiments of this disclosure, the following steps can be used to synthesize paired training image data:
[0066] i. Crop blank background areas without text regions from scene images or document images to create a background image database.
[0067] ii. Collect a corpus to use as a text content database.
[0068] iii. Randomly select a background image Tb from the background image database, where Tb represents the real background image corresponding to the style image.
[0069] iv. Randomly select two lines of text s and t from the text content database, as the first text s and the second text t.
[0070] v. Randomly generate style rendering parameters such as text coordinates, font size, and font type.
[0071] vi. In the binary solid color background, apply the style rendering parameters generated in step v to the two lines of text selected in step iv to obtain the real text mask image T corresponding to the style image. ms and the corresponding real text mask image T of the converted content image mt .
[0072] vii. In a pure gray background image, apply the SimSun font type (or other font type) to the second text t selected in step iv to obtain the above content image I. t viii. In the pure gray background image, the style rendering parameters generated in step v are applied to the second text t selected in step iv to obtain the actual converted content image T corresponding to the converted content image. t .
[0073] ix. In the background image selected in step iii, the style rendering parameters generated in step v are applied to the first text s selected in step iv to obtain the aforementioned style image I. s .
[0074] x. In the background image selected in step iii, the style rendering parameters generated in step v are applied to the second text t selected in step iv to obtain the real target image T corresponding to the target image. f.
[0075] At box S104, a background image and a corresponding first text mask image are obtained from the style image by the background erasure module. The background image and the first text mask image are obtained in parallel. The first text mask image can be a binary matrix. The first text mask image can have a black background and be white at the location of the first text in the corresponding style image. In some embodiments of this disclosure, the first text mask image is used to calculate the loss function of the background erasure module.
[0076] In some embodiments of this disclosure, the background erasure module may include a first generative adversarial network (GAN). The first generator in the first GAN employs a UNet-shaped encoder-decoder structure. Figure 2 In one example, the first generator may include a first encoder 211, a first decoder 212, and a second decoder 213. The first decoder 212 and the second decoder 213 are respectively symmetrical to the first encoder 211. In some embodiments of this disclosure, the first encoder 211 processes the style image I... s Encode to obtain style image I s The latent space representation (which can also be considered as the feature representation of the style image). Style image I s The latent space representation is input to the first decoder 212 and the second decoder 213. The first decoder 212 reconstructs the latent space representation into the background image O. b The second decoder 213 reconstructs the latent space representation into the first text mask image O. ms .
[0077] In some embodiments of this disclosure, the first encoder 211 may consist of a 3-layer convolutional neural network with a downsampling stride of 2. The first decoder 212 and the second decoder 213 may each consist of a 3-layer convolutional neural network with an upsampling stride of 2. Each convolutional neural network is followed by a batch regularization layer and a leaky ReLU activation function. Since using transposed convolution to pad with zeros results in fixed pixel values, the generated background image will contain grid-like texture noise at fixed intervals. Embodiments of this disclosure employ bilinear interpolation to achieve feature upsampling. In some embodiments of this disclosure, skip connections may be added between the feature layers of the encoder and decoder with mirror dimensions. This connection helps the encoder recover the background texture information lost during downsampling.
[0078] In some embodiments of this disclosure, the first discriminator in the first generative adversarial network can draw inspiration from PatchGAN and consists of four downsampled convolutional blocks. Each convolutional block includes a convolutional layer, a batch regularization layer, and a leaky ReLU activation function, and finally outputs downsampled image patches to determine the authenticity of image details.
[0079] To make the texture details in the generated background image closer to the style image, in some embodiments of this disclosure, three loss functions are used to supervise the background erasure module. The loss function of the background erasure module is equal to the sum of the first loss function, the second loss function, and the third loss function.
[0080] The first loss function is applied to the generated background image O. b The maximum entropy loss function. The first loss function can be expressed as:
[0081]
[0082] Among them, D B I represents the first discriminator in the first generative adversarial network used to determine the authenticity of an image. s Represents style image, O b T represents the background image. b This represents the real background image corresponding to the style image. The real background image T corresponding to the style image. b It can be obtained from the training dataset. The first loss function L... GAN1 By utilizing the concept of adversarial generation, the generator is guided to generate more deceptive images, while the discriminator is guided to better distinguish between real and generated images.
[0083] The second loss function is the Dice loss function for the first text mask image. The second loss function can be expressed as:
[0084]
[0085] Among them, O ms T represents the first text mask image. ms This represents the real text mask image corresponding to the style image, where N represents the size of the first text mask image. O ms The i-th element in. T represents ms The i-th element in the image. The real text mask image T corresponding to the style image. ms It can be obtained from the training dataset. The second loss function L Dice1 The generated first text mask image O was supervised. ms and the real text mask image T corresponding to the style imagems The similarity between them.
[0086] The third loss function is the L1 loss function for the background image. The third loss function can be expressed as:
[0087] L L1_1 =||T b -O b ||1 (3)
[0088] Among them, O b T represents the background image. b This represents the real background image corresponding to the style image. The loss function of the background erasure module can be expressed as the sum of the three loss functions mentioned above:
[0089] L B =L GAN1 +L Dice1 +L L1_1 (4)
[0090] The background erasure module has two branches, which generate a background image and a corresponding first text mask image through two independent decoders, thereby supervising the generation of a cleaner background image from both the perspectives of background texture restoration and text region positioning.
[0091] At box S106, the text conversion module converts the second text style attribute of the second text in the content image into a first text style attribute to obtain a converted content image and a second text mask image corresponding to the converted content image. The converted content image and the second text mask image are obtained in parallel. The second text mask image can be a binary matrix. The second text mask image can have a black background and be white at the positions of the second text in the corresponding content image. In some embodiments of this disclosure, the second text mask image is used to calculate the loss function of the text conversion module.
[0092] In some embodiments of this disclosure, the text conversion module may include an encoder-decoder structure based on a fully convolutional network. Figure 2 In the example, the encoder-decoder structure may include: a second encoder 221, a third encoder 222, a third decoder 223, and a fourth decoder 224.
[0093] In some embodiments of this disclosure, the second encoder 221 can process the style image I s Encoding is performed to obtain style image features. The third encoder 222 encodes the content image I. tEncoding is performed to obtain content image features. Then, the style image features and content image features are concatenated into a fused feature block along the channel dimension. Next, the third decoder 223 reconstructs the transformed content image O from the fused feature block. t The fourth decoder 224 can reconstruct the fused feature blocks into a second text mask image O. mt .
[0094] In some embodiments of this disclosure, the second encoder 221 and the third encoder 222 may each employ three downsampled convolutional layers and four residual convolutional coding blocks. The third decoder 223 and the fourth decoder 224 may each employ three bilinear interpolated upsampled convolutional layers, wherein each upsampled layer is followed by a batch regularization layer and a leaky ReLU activation function.
[0095] In some embodiments of this disclosure, such as Figure 2 As shown, to better generate text outlines, a pre-trained text recognition model 240 is introduced during the training of the image generation model to supervise the readability of the text in the converted content image. The text recognition model 240 can predict the converted content image, and the prediction result of the text recognition model 240 is added to the calculation of the loss function of the text conversion module 220. In one example, the text recognition model 240 can output the predicted text and the predicted probability of each character in the predicted text for the converted content image. Specifically, the text recognition model 240 may include a dictionary. The text recognition model 240 can calculate the predicted probability that each character in the predicted text is a character in the dictionary.
[0096] In some embodiments of this disclosure, the text recognition model 240 may employ a TRBA (TPS-ResNetBiLSTM-Attention) configuration. In this configuration, a spatial transformation network based on Thin-Plate Spline (TPS) interpolation is used; a ResNet network is used to extract image features; BiLSTM is used for sequence modeling; and attention is used for sequence prediction. The text recognition model 240 can be pre-trained on a combination of the publicly available Chinese and English datasets MJSynth and SynthText.
[0097] To ensure that the text style in the converted content image closely resembles the text style in the style image, in some embodiments of this disclosure, three loss functions are used to supervise the text conversion module. The loss function of the text conversion module is equal to the sum of the fourth, fifth, and sixth loss functions.
[0098] The fourth loss function is the multi-class cross-entropy loss function applied to the second text in the transformed content image. The fourth loss function can be expressed as:
[0099]
[0100] Where M represents the number of characters in the dictionary used by the text recognition model, k represents the number of characters in the second text in the converted content image, and p ic y represents the predicted probability that the i-th character in the second text of the transformed content image is the c-th character in the dictionary. ic This indicates whether the i-th character in the second text of the converted content image is the c-th character in the dictionary. If y ic A value of 1 indicates that the i-th character in the second text of the converted content image is the c-th character in the dictionary. If y ic A value of 0 indicates that the i-th character in the second text of the converted content image is not the c-th character in the dictionary. In some embodiments of this disclosure, y ic The value can be input by the user. Fourth loss function L reco1 By simulating the human perception process through a text recognition model, the reasonableness of the shape and outline of the second text in the converted content image is supervised.
[0101] The fifth loss function is the Dice loss function for the second text mask image. The fifth loss function can be expressed as:
[0102]
[0103] Among them, O mt T represents the second text mask image. mt This represents the real text mask image corresponding to the converted content image, where n represents the size of the second text mask image. The real text mask image T corresponding to the converted content image. mt It can be obtained from the training dataset.
[0104] The sixth loss function is the L1 loss function for the transformed content image. The sixth loss function can be expressed as:
[0105] L L1_2 =||T t -O t ||1 (7)
[0106] Among them, O t T represents the converted content image. t This represents the actual converted content image corresponding to the converted content image. The actual converted content image T corresponds to the converted content image. t It can be obtained from the training dataset.
[0107] The loss function of the text conversion module can be expressed as the sum of the three loss functions mentioned above:
[0108] L T =L reco1 +L Dice2 +L L1_2 (8)
[0109] During training, the text conversion module is supervised by the text recognition model, which can improve the readability of the generated images.
[0110] At box S108, the converted content image and the background image are fused by the fusion module to obtain the target image. In some embodiments of this disclosure, the fusion module may include: a second generative adversarial network (GAN). Figure 2 (Not shown in the image). The second generator in the second generative adversarial network may include a first convolutional layer and a second convolutional layer. The first convolutional layer may be a 1×1 convolutional layer. The second convolutional layer may be a ResNet Block convolutional layer.
[0111] In some embodiments of this disclosure, the background image and the converted content image can be stitched together into a stitched image along the channel dimension. Embodiments of this disclosure do not limit the stitching order of the background image and the converted content image. Then, the dimensionality of the stitched image can be reduced through a first convolutional layer to perform preliminary image fusion. The fused image features can be obtained in this step. Next, a second convolutional layer interacts with the feature information on the fused image features to generate the target image.
[0112] In some embodiments of this disclosure, such as Figure 2 As shown, to better generate text outlines, a pre-trained text recognition model 240 is introduced during the training of the image generation model to supervise the readability of the text in the target image. The text recognition model 240 can predict the target image, and the prediction results are added to the calculation of the loss function of the fusion module 230. In one example, the text recognition model 240 can output the predicted text and the predicted probability of each character in the predicted text for the target image. Specifically, the text recognition model 240 may include a dictionary. The text recognition model 240 can calculate the predicted probability that each character in the predicted text is a character in the dictionary.
[0113] exist Figure 2 The image shows a text recognition model 240 that can receive a second text mask image O from a text conversion module 220. mt and the target image O from the fusion module 230 f However, the second text mask image O mt and the target image O from the fusion module 230 fThe text recognition model 240 is not input simultaneously. The text recognition model 240 is based on the second text mask image O. mt and target image O f The loss functions of the text conversion module 220 and the fusion module 230 are calculated respectively. In other words, the text conversion module 220 and the fusion module 230 only reuse the text recognition model 240 in a time-sharing manner, rather than using the text recognition model 240 simultaneously.
[0114] To make the target image closer to the real image, in some embodiments of this disclosure, five loss functions are used to supervise the fusion module. The loss function of the fusion module is equal to the sum of the seventh, eighth, ninth, tenth, and eleventh loss functions.
[0115] The seventh loss function is the cross-entropy loss function applied to the second text in the target image. The seventh loss function can be expressed as:
[0116]
[0117] Where M represents the number of characters in the dictionary used by the text recognition model, k represents the number of characters in the second text in the target image, and P ic Y represents the predicted probability that the i-th character in the second text of the target image is the c-th character in the dictionary. ic Indicate whether the i-th character in the second text of the target image is the c-th character in the dictionary. If Y ic A value of 1 indicates that the i-th character in the second text of the target image is the c-th character in the dictionary. If Y... ic A value of 0 indicates that the i-th character in the second text of the target image is not the c-th character in the dictionary. In some embodiments of this disclosure, Y ic The value can be entered by the user. Seventh loss function L reco2 By simulating the human perception process through a text recognition model, the reasonableness of the shape and outline of the second text in the target image is supervised, and the recognizability of the second text in the target image is measured.
[0118] The eighth loss function is a perceptual loss function for the target image. In some embodiments of this disclosure, the fusion module 230 may include an image feature extraction model. In method 100, the image feature extraction model can be used to extract features from the target image and the corresponding real target image. In one example, the image feature extraction model may be a VGG-19 classification model pre-trained on the ImageNet image classification task. The VGG-19 classification model pre-trained on the ImageNet image classification task can be used to extract features from the target image and the corresponding real target image, respectively, sequentially extracting the feature maps output by the five ReLU layers in the VGG-19 model, and then calculating the L1 loss of the corresponding feature maps. The eighth loss function can be expressed as:
[0119]
[0120] Where, φ i (·) represents the i-th feature map output by the image feature extraction model, O f T represents the target image. f Let T represent the real target image corresponding to the target image, and h represent the number of feature maps output by the image feature extraction model. f It can be obtained from the training dataset. The perceptual loss supervises the reconstruction of image style through the image feature extraction model, which can guide the fusion module to generate better text shapes.
[0121] The ninth loss function is a texture loss function for the target image. The ninth loss function can be expressed as:
[0122]
[0123] Among them, G i (·) represents the Gram matrix of the i-th feature map output by the image feature extraction model, i.e., G i =φ i φ i T , where φ i This represents the i-th feature map. O f T represents the target image. f M represents the real target image corresponding to the target image, h represents the number of feature maps output by the image feature extraction model, and M represents the number of feature maps. i This represents the size of the i-th feature map. In the VGG-19 model, h can be equal to 5. The real target image T corresponds to the target image. f It can be obtained from the training dataset. The ninth loss function L tex It can guide the fusion module to generate more realistic texture details.
[0124] The tenth loss function is the maximum entropy loss function for the target image. The tenth loss function can be expressed as:
[0125]
[0126] Among them, D F I represents the second discriminator in the second generative adversarial network. t Represents the content image, O f T represents the target image. f This represents the real target image corresponding to the target image. The real target image T corresponding to the target image. f It can be obtained from the training dataset.
[0127] The eleventh loss function is the L1 loss function for the target image. The eleventh loss function can be expressed as:
[0128] L L1_3 =||T f -O f ||1 (13)
[0129] Among them, O f T represents the target image. f This represents the real target image corresponding to the target image. The real target image T corresponding to the target image. f It can be obtained from the training dataset.
[0130] The loss function of the fusion module can be expressed as the sum of the five loss functions mentioned above:
[0131] L F =L reco2 +L per +L tex +L GAN2 +L L1_3 (14)
[0132] At box S110, the loss functions of the background erasure module, the text conversion module, and the fusion module are minimized to update the parameters of the image generation model. The loss function of the image generation model is equal to the sum of the loss functions of the background erasure module, the text conversion module, and the fusion module. In some embodiments of this disclosure, the loss function L of the background erasure module can be minimized separately during training. B The loss function L of the text conversion module T The loss function L of the fusion module F By minimizing L B L T and L FThe parameters of the image generation model can be optimized. The optimized parameters can then be used to update the parameters of the image generation model.
[0133] In the above process, the training operations performed at boxes S104 to S110 are operations performed in one round of training. These training operations can be performed in multiple rounds to obtain better training results. In some embodiments of this disclosure, the training operations can end after a predetermined number of training rounds are completed, or they can end when the loss function of the image generation model converges.
[0134] Figure 3 This diagram illustrates three sets of style images, content images, and the target image derived from them. Figure 3 In the example, the background image of style image 31a can be obtained through the background erasure module, and the text content of content image 31b can be converted into the text style of style image 31a through the text conversion module. Then, the background image of style image 31a and the text information in the converted content image are merged through the fusion module to obtain target image 31c. Figure 3 As shown, target image 31c is equivalent to converting the background and text style of content image 31b according to style image 31a. Similarly, target image 32c is equivalent to converting the background and text style of content image 32b according to style image 32a. Target image 33c is equivalent to converting the background and text style of content image 33b according to style image 33a.
[0135] In some embodiments of this disclosure, the image generation model may further include a text-to-image module. The text-to-image module can apply a font type (e.g., SimSun font type, or other font type) to the input target text, and then add the target text to a plain gray background image to obtain a content image. In actual use of the trained image generation model to generate images, a style image and target text can be input to the image generation model. The image generation model can first convert the target text into a content image, and then generate a target image with a background image from the style image and the target text based on the content image and the style image. The text style of the target text is the same as the text style of the style image.
[0136] Figure 4 A schematic block diagram of an apparatus 400 for training an image generation model according to an embodiment of the present disclosure is shown. Figure 4 As shown, the device 400 may include a processor 410 and a memory 420 storing a computer program. When the computer program is executed by the processor 410, the device 400 is made capable of performing actions such as... Figure 1The steps of method 100 are shown below. In one example, device 400 may be a computer device or a cloud computing node. Device 400 may acquire a training image set. The training image set includes: a style image containing first text and a content image containing second text. The first text has a first text style attribute. The second text has a second text style attribute. Device 400 may obtain a background image and a corresponding first text mask image from the style image via a background erasure module. The first text mask image is used to calculate the loss function of the background erasure module. Device 400 may convert the second text style attribute of the second text in the content image to a first text style attribute via a text conversion module to obtain a converted content image and a corresponding second text mask image. The second text mask image is used to calculate the loss function of the text conversion module. Device 400 may fuse the converted content image with the background image via a fusion module to obtain a target image. Device 400 may minimize the loss functions of the background erasure module, the text conversion module, and the fusion module to update the parameters of the image generation model.
[0137] In some embodiments of this disclosure, the background erasure module may include a first generative adversarial network (GAN). The first generator in the first GAN employs a UNet-shaped encoder-decoder structure. The apparatus 400 may encode a style image using a first encoder in the first generator to obtain a latent space representation of the style image. The apparatus 400 may reconstruct a background image from the latent space representation using a first decoder symmetrical to the first encoder in the first generator. The apparatus 400 may reconstruct a first text mask image from the latent space representation using a second decoder symmetrical to the first encoder in the first generator.
[0138] In some embodiments of this disclosure, the device 400 can supervise the text readability of the converted content image through a pre-trained text recognition model.
[0139] In some embodiments of this disclosure, the text conversion module may include an encoder-decoder structure based on a fully convolutional network. The device 400 may use a second encoder in the encoder-decoder structure to encode a style image to obtain style image features. The device 400 may use a third encoder in the encoder-decoder structure to encode a content image to obtain content image features. The device 400 may concatenate the style image features and content image features into a fused feature block along the channel dimension. The device 400 may use the third decoder in the encoder-decoder structure to reconstruct the fused feature block into a converted content image. The device 400 may use a fourth decoder in the encoder-decoder structure to reconstruct the fused feature block into a second text mask image.
[0140] In some embodiments of this disclosure, the fusion module may include a second generative adversarial network (GAN). The second generator in the second GAN includes a first convolutional layer and a second convolutional layer. The device 400 can stitch a background image and a transformed content image into a stitched image along the channel dimension. The device 400 can reduce the dimensionality of the stitched image and obtain fused image features through the first convolutional layer. The device 400 can interact with feature information on the fused image features through the second convolutional layer to generate a target image.
[0141] In some embodiments of this disclosure, the device 400 can supervise the text readability of a target image through a pre-trained text recognition model.
[0142] In some embodiments of this disclosure, the training image set may further include a real target image corresponding to the target image. The apparatus 400 may use an image feature extraction model to extract features from the target image and the real target image corresponding to the target image.
[0143] In embodiments of this disclosure, processor 410 may be, for example, a central processing unit (CPU), a microprocessor, a digital signal processor (DSP), a processor based on a multi-core processor architecture, etc. Memory 420 may be any type of memory implemented using data storage technologies, including but not limited to random access memory, read-only memory, semiconductor-based memory, flash memory, disk storage, etc.
[0144] Furthermore, in embodiments of this disclosure, the device 400 may also include an input device 430, such as a keyboard or mouse, for inputting a training image set. Additionally, the device 400 may also include an output device 440, such as a display, for outputting the target image and parameters of the image generation model.
[0145] In other embodiments of this disclosure, a computer-readable storage medium storing a computer program is also provided, wherein the computer program, when executed by a processor, is capable of performing the following functions: Figure 1 The steps of the method shown.
[0146] In summary, the method and apparatus for training an image generation model according to embodiments of this disclosure supervise the generation of a background image and a corresponding first text mask image through two decoders in the background erasure module. Through positive background supervision and negative text supervision, the background texture of the text region edges in the style image can be reconstructed more effectively. In the text conversion module, the generation of the converted content image and the corresponding second text mask image is supervised by two decoders. Through positive text supervision and negative background supervision, the text style of the background region edges in the content image can be reconstructed more effectively. In the fusion module, the background image and the converted content image are fused using a channel concatenation dimensionality reduction method to generate a target image. Furthermore, embodiments of this disclosure simulate the human process of recognizing text using a pre-trained text recognition model, predicting the converted content image and the target image, making the text in the generated image more readable. The image generation model constructed by embodiments of this disclosure can generate relatively realistic synthetic images. By editing the text content in few-sample or long-tailed distributed data, the image generation model can synthesize new training data, which can be applied to the field of data augmentation, especially data augmentation of rare characters in optical character recognition tasks.
[0147] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of apparatuses and methods according to various embodiments of the present disclosure. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of an instruction containing one or more executable instructions for implementing a specified logical function. In some alternative implementations, the functions marked in the blocks may occur in a different order than those marked in the drawings. For example, two consecutive blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, may be implemented using a dedicated hardware-based system that performs the specified function or action, or using a combination of dedicated hardware and computer instructions.
[0148] Unless otherwise expressly indicated by the context, the singular form of words used herein and in the appended claims includes the plural form, and vice versa. Thus, when referring to the singular, the plural form of the corresponding term is generally included. Similarly, the terms “comprising” and “including” shall be interpreted as including rather than exclusively. Likewise, the terms “including” and “or” shall be interpreted as including unless such interpretation is expressly prohibited herein. Where the term “example” is used herein, particularly when it follows a set of terms, the “example” is merely exemplary and illustrative and should not be considered exclusive or extensive.
[0149] Further aspects and scope of adaptation become apparent from the description provided herein. It should be understood that various aspects of this application may be implemented individually or in combination with one or more other aspects. It should also be understood that the descriptions and specific embodiments herein are for illustrative purposes only and are not intended to limit the scope of this application.
[0150] Several embodiments of this disclosure have been described in detail above. However, it is obvious that those skilled in the art can make various modifications and variations to the embodiments of this disclosure without departing from the spirit and scope of this disclosure. The scope of protection of this disclosure is defined by the appended claims.
Claims
1. A method for training an image generation model, the image generation model comprising a background erasure module, a text conversion module, and a fusion module, the method comprising: Obtain a training image set, wherein the training image set includes: a style image containing a first text and a content image containing a second text, the first text having a first text style attribute and the second text having a second text style attribute; The background erasure module obtains a background image and a first text mask image corresponding to the background image from the style image, wherein the first text mask image is used to calculate the loss function of the background erasure module; The text conversion module converts the second text style attribute of the second text in the content image into the first text style attribute to obtain the converted content image and the second text mask image corresponding to the converted content image. The second text mask image is used to calculate the loss function of the text conversion module. The conversion of the content image and the background image are fused together by the fusion module to obtain the target image; and The loss functions of the background erasure module, the text conversion module, and the fusion module are minimized to update the parameters of the image generation model; The loss function of the background erasure module is equal to the sum of the first loss function, the second loss function and the third loss function. The first loss function is the maximum entropy loss function for the generated background image, the second loss function is the Dice loss function for the first text mask image, and the third loss function is the L1 loss function for the background image. The loss function of the text conversion module is equal to the sum of the fourth loss function, the fifth loss function and the sixth loss function. The fourth loss function is the multi-class cross-entropy loss function for the second text in the converted content image, the fifth loss function is the Dice loss function for the mask image of the second text, and the sixth loss function is the L1 loss function for the converted content image. The loss function of the fusion module is equal to the sum of the seventh, eighth, ninth, tenth and eleventh loss functions. The seventh loss function is the cross-entropy loss function for the second text in the target image, the eighth loss function is the perceptual loss function for the target image, the ninth loss function is the texture loss function for the target image, the tenth loss function is the maximum entropy loss function for the target image, and the eleventh loss function is the L1 loss function for the target image.
2. The method according to claim 1, wherein, The background erasure module includes: a first generative adversarial network, wherein the first generator in the first generative adversarial network adopts a UNet-shaped encoder-decoder structure, wherein obtaining a background image and a first text mask image corresponding to the background image from the style image through the background erasure module includes: The style image is encoded using the first encoder in the first generator to obtain the latent space representation of the style image; The latent space representation is reconstructed into the background image using a first decoder symmetrical to the first encoder in the first generator; and The latent space representation is reconstructed into the first text mask image using the second decoder, which is symmetrical to the first encoder in the first generator.
3. The method according to claim 2, wherein, The training image set also includes: a real background image corresponding to the style image, and a real text mask image corresponding to the style image; the loss function of the background erasure module is equal to the sum of the first loss function, the second loss function, and the third loss function; The first loss function is expressed as: + in, This represents the first discriminator in the first generative adversarial network. This refers to the style image. This refers to the background image. This represents the real background image corresponding to the style image; The second loss function is expressed as: in, Represents the first text mask image The first text mask image represents the real text mask image corresponding to the style image, and N represents the size of the first text mask image. The third loss function is expressed as: in, This refers to the background image. This represents the real background image corresponding to the style image.
4. The method according to claim 1, further comprising: The readability of the text in the converted content image is supervised by a pre-trained text recognition model.
5. The method according to claim 4, wherein, The text conversion module includes: an encoder-decoder structure based on a fully convolutional network, which converts the second text style attribute of the second text in the content image into the first text style attribute to obtain a converted content image and a second text mask image corresponding to the converted content image. The style image is encoded using the second encoder in the encoder-decoder structure to obtain style image features; The content image is encoded using the third encoder in the encoder-decoder structure to obtain content image features; The style image features and the content image features are concatenated into a fused feature block along the channel dimension; The fused feature blocks are reconstructed into the transformed content image using the third decoder in the encoder-decoder structure; and The fused feature block is reconstructed into the second text mask image using the fourth decoder in the encoder-decoder structure.
6. The method according to claim 5, wherein, The training image set also includes: a real text mask image corresponding to the converted content image, and a real converted content image corresponding to the converted content image. The loss function of the text conversion module is equal to the sum of the fourth loss function, the fifth loss function, and the sixth loss function. The fourth loss function is expressed as: Where M represents the number of characters in the dictionary used by the text recognition model, and k represents the number of characters in the second text in the converted content image. This represents the predicted probability that the i-th character in the second text of the converted content image is the c-th character in the dictionary. This indicates whether the i-th character in the second text of the converted content image is the c-th character in the dictionary. If... A value of 1 indicates that the i-th character in the second text of the converted content image is the c-th character in the dictionary; The fifth loss function is expressed as: in, Represents the second text mask image This represents the real text mask image corresponding to the converted content image, where n represents the size of the second text mask image; The sixth loss function is expressed as: in, This represents the converted content image. This represents the actual converted content image corresponding to the converted content image.
7. The method according to claim 1, wherein, The fusion module includes a second generative adversarial network (GAN), wherein a second generator in the GAN includes a first convolutional layer and a second convolutional layer. The fusion module fuses the transformed content image with the background image to obtain a target image, comprising: The background image and the converted content image are stitched together into a stitched image along the channel dimension; The first convolutional layer reduces the dimensionality of the stitched image and obtains the fused image features; and The second convolutional layer interacts with the feature information on the fused image features to generate the target image.
8. The method according to claim 7, further comprising: The readability of the text in the target image is supervised by a pre-trained text recognition model.
9. The method according to claim 8, wherein, The training image set further includes: a real target image corresponding to the target image; the method further includes: using an image feature extraction model to extract features from the target image and the real target image corresponding to the target image; The loss function of the fusion module is equal to the sum of the seventh loss function, the eighth loss function, the ninth loss function, the tenth loss function, and the eleventh loss function. The seventh loss function is expressed as: Where M represents the number of characters in the dictionary used by the text recognition model, and k represents the number of characters in the second text in the target image. This represents the predicted probability that the i-th character in the second text of the target image is the c-th character in the dictionary. Indicate whether the i-th character in the second text of the target image is the c-th character in the dictionary. If A value of 1 indicates that the i-th character in the second text of the target image is the c-th character in the dictionary; The eighth loss function is expressed as: in, This represents the i-th feature map output by the image feature extraction model. Represents the target image The image represents the real target image corresponding to the target image, and h represents the number of feature maps output by the image feature extraction model. The ninth loss function is expressed as: in, This represents the Gram matrix of the i-th feature map output by the image feature extraction model. Represents the target image Let h represent the real target image corresponding to the target image, and let h represent the number of feature maps output by the image feature extraction model. This represents the size of the i-th feature map; The tenth loss function is expressed as: + in, This represents the second discriminator in the second generative adversarial network. This refers to the content image. Represents the target image This represents the real target image corresponding to the target image; The eleventh loss function is expressed as: in, Represents the target image This represents the real target image corresponding to the target image.
10. An apparatus for training an image generation model, comprising: At least one processor; as well as At least one memory storing a computer program; When the computer program is executed by the at least one processor, the apparatus performs the steps of the method according to any one of claims 1 to 9.