Defect Transfer GAN: Diverse Defect Synthesis for Data Augmentation

1 INTRODUCTION . Automated Visual Inspection ( AVI ) is vital for quality control in modern production lines . Despite the fact that AVI has been studied for decades , it remains a challenging task with many open research questions await to be answered . One of the main challenges in data-driven AVI is the acquisition of suitable training data . This is for two reasons : First , collecting a vast amount of labelled data is usually labor-intensive and time-consuming . In many cases , even experts are required to identify where and what to look for . However , the acquired label information is task-specific and can not be reused or transferred to a new task in most cases . Thus , the tedious labelling process must be repeated for each new product , even if its defect is similar to other products in people ’ s eyes . Second , in real-world scenarios such as highly optimized production lines , a more severe problem emerges : data imbalance . Only very few defective parts are produced by design . Moreover , the acquired anomaly images from a single product are lacking diversity and may not capture the full defect distribution . Training a robust deep neural network model in such conditions is very challenging . Since collecting sufficient real-world defective samples is impractical , algorithms to synthesize required images became a focus in research . Image synthesis through Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) has shown promising performance in recent years . But it also requires large amounts of balanced data which are not available in most industrial use cases , in particular for irregular defect patterns and large variation . Therefore , GANs tend to overfit to the training examples when trained with little data ( Karras et al. , 2020a ) . In this work , we tackle these issues by exploiting cross-domain information : we first define two sets of domains—foreground domains and background domains . The foreground domain describes a set of images that contains a specific foreground content to be grouped into a distinctive category , and each content has a different style . The background domain instead is considered as a group of images that shares similar structural appearance over the whole image . For example , we can set foreground domains as defect types and background domains as product types while the styles of defects indicate their artistic looks such as light or heavy strokes . Building upon StarGAN v2 ( Choi et al. , 2020 ) , the concept underlying this work is to transfer and generate foreground contents with a variety of styles across different background domains , as illustrated in Figure 1 . The contributions of this work are three-fold : First , we introduce Defect Transfer GAN ( DT-GAN ) , a model that learns transferring existing foreground content and generating novel contents onto different backgrounds at the same time . In the real-world scenario , it allows defect inspection networks to learn from a variety of synthetic defective images by composing the foreground defects together with various non-defective images from different products . Second , DT-GAN is able to disentangle the foreground defect-specific content and the defect-irrelevant background in a weakly-supervised manner . Third , extensive experiments show that our method can generate diverse and real-looking defective samples even for products with only 20 real defective images . These defective images generated by DT-GAN boost the performance in defect inspection networks significantly . 2 RELATED WORK . GANs have shown their power in many computer vision tasks such as image synthesis ( Lučić et al. , 2019 ) , style translation ( Johnson et al. , 2016 ) , super-resolution ( Ledig et al. , 2017 ) , image impainting ( Pathak et al. , 2016 ) and many other applications . To quantify the performance of GANs , visual quality and the diversity of generated images are considered as two of the most important criteria . Recent models address these requirements either by dedicated loss functions ( Mao et al. , 2019b ; Yang et al. , 2019 ) or architectural design ( Brock et al. , 2019 ) . StyleGAN v2 ( Karras et al. , 2020b ) , the latest state-of-the-art model in image synthesis , introduces stochastic variation in image generating process by adding per-pixel noise after each convolution . However , it is non-trivial to adapt the model to transform given input images due to the design of the generator . In contrast , image-to-image translation methods ( Isola et al. , 2017 ) provide a way to recover the connection between inputs and the generated images while encouraging diversity . For example , Zhu et al . ( 2017b ) and Huang et al . ( 2018 ) impose consistent mappings in latent space to achieve the goal . Some approaches ( Ma et al. , 2019 ; Park et al. , 2019 ) use reference images as guidance to generate diverse outputs . Mokady et al . ( 2020 ) further extends the translation task from styles to contents . It learns to identify a specific content in a given input ( e.g. , a specific pair of glasses ) and transfer it to the target image . However , aforementioned methods only consider the translation between two domains and their extension to multiple domains is non-trivial . Surface defect detection is one of the important tasks in real-world industrial manufacturing . It aims at identifying and classifying defects with the help of machine vision . Traditional methods ( Ngan et al. , 2011 ) build models upon hand-crafted feature extractors , which are unstable and outperformed by deep learning based models . However , the performance and generalization ability of deep learning approaches are restricted due to limited number of defective samples in real-world scenarios . Data augmentation aims to enrich the training dataset by introducing different kinds of invariance for the model to capture . Several recent works ( Niu et al. , 2020 ; Zhang et al. , 2021 ) have proposed to adopt GANs as a data augmentation method to generate realistic defective samples . Among them , Defect-GAN ( Zhang et al. , 2021 ) tries to capture the stochastic variation within defects by mimicking the defacement and restoration processes . However , it still learns a deterministic mapping between inputs and outputs while DT-GAN achieves multi-modality by varying styles . Moreover , our method can generate realistic defects with sophisticated patterns copied from real-world defective samples . 3 METHODOLOGY . Our primary aim is to perform unpaired image-to-image translation across multiple foreground domains within a single model . In our use case , the foreground domains refer to the defect types , which means we want to achieve translations between different types of defects while the background remains unaffected . We assume that there is always an adequate amount of normal samples ( e.g. , non-defective ) available , while anomaly samples are rare and hard to acquire . 3.1 PROPOSED FRAMEWORK . Our framework builds on StarGAN v2 , a multimodal image-to-image translation model . Given an input image x ∈ X and an arbitrary domain y ∈ Y , StarGAN v2 generates a domain specific style code in a learned style space and outputs an image that is stylized to fit the domain of y . Its network architecture consists of four modules : a generator , a mapping network , a style encoder and a discriminator . We modify and extend all four modules ( see Figure 2 ) and describe the key differences in details as below . Style-Content Separation . Given a latent code z and a domain y , the mapping network M ( Figure 2 ( b ) ) generates a style code s = My ( z ) and a domain specific content c = My ( z ) in different branches . It is worth mentioning that My here denotes an output of M corresponding to the domain y . This feature allows our method to separately model the structural appearance ( i.e . content ) and its artistic looks ( i.e . style ) , which is essential because applying different styles to the same content enriches the diversity of outputs . By randomly sampling z from a standard normal distribution and y from all available foreground domains , M is able to produce diverse style codes and domain specific contents . The encoder E ( Figure 2 ( c ) ) extracts the style code s = Ey ( x ) and the domain specific content c = Ey ( x ) from an given image x , which reflect the characteristics of reference images instead of randomly sampled noise . Foreground/Background ( FG/BG ) Disentanglement . The generator G ( Figure 2 ( a ) ) translates an input image x into an output image G ( x , s̃ , c̃ ) according to given domain specific style code s̃ and content c̃ , which are provided either by the mapping network M when generating from random noise or by the style-content encoder E when transferring an existing content from a reference image . To achieve a FG/BG disentanglement , we split the channels of the three-dimensional feature map ( i.e. , H × W × C ) at the bottle neck of G into two parts . The model is then forced to encode the background into the first channels and the domain specific content ĉ into the latter channels by classification losses as discussed in Section 3.2. ĉ can then be replaced with content c̃ from the target domain . The adaptive instance normalization ( AdaIN ) ( Huang & Belongie , 2017 ) is then used to inject s̃ into c̃ during the decoding process while the background BGG ( x ) is decoded separately . StarGAN v2 learns FG and BG together which leads to a conditional relationship between both . Our disentanglement and separate encoding break this conditioning and therefore enable our method to freely combine FG and BG as well as learn the full variation of FG content . Finally , BGG ( x ) and c̃ are concatenated together and then fused before output . Multi-task discriminator with auxiliary classifiers . The discriminator D ( Figure 2 ( d ) ) is a multitask discriminator with two auxiliary classifiers : a foreground domain classifier and a background domain classifier . This feature strengthens the disentanglement of FG and BG by first ensuring the input image x contains a domain specific content that can be recognized by the foreground domain classifier independent of the background . Later , each branch Dy in the multi-task discriminator D is trained to determine if an image x is a real image of its foreground domain or a fake image G ( x , s , c ) generated by G. Apart from that , one extra branch BGcls is attached to decide whether the background information of the input images is well preserved . Content Transfer . Mokady et al . ( 2020 ) introduced a concept that a model should be able to identity the difference between two domains when one of the domains contains a feature that the other does not have . We refer to this concept as ‘ anchor ’ and extend to multiple domains ( > 2 ) by the FG/BG disentanglement , the multi-task discriminator and the foreground content classifier in D. We treat domain Normal as the anchor domain i.e . set the domain specific content to zero , because a normal image has no domain specific content in our definition . As a result , we can now transfer contents between all combination of FG and BG domains ( see Figure 10 ) . Compared to StarGAN v2 , our method not only models style codes and contents separately but also disentangles the foreground and background of an image in a weakly-supervised manner . These features allow explicit control over output images by combining desired style codes and contents from one of the subnetworks with the input images . Therefore , it leads to higher variance regarding the location , structural pattern and artistic style of defects in the synthetic images of DT-GAN .

Summary:  The paper prposes a StarGAN based model to distangle the defect foreground, background, transfer the style of foreground and then synthesize the defected image for different products. The quality of the synthesized defect images are evaluated using FID, KID. The approach is also appled to augment defect images to improve the performance of defect classification.