[Paper review] Few-Shot Defect Image Generation via Defect-Aware Feature Manipulation

0

Paper Seminar

목록 보기
10/10

Introduction

  • 제조 현장에서 진행되는 Defect inspection은 결함 검출, 분류, Localization 등을 주요 업무로 함
    지금까지 자동화된 defect inspection system에 대한 연구가 진행되어 왔음
    그러나, 실제 결함 이미지의 수가 적고 높은 수집 비용으로 인해 다양한 결함 이미지를 확보하는 것이 어려우며 이러한 문제를 Data insuffciency problem 이라고 함
    따라서 일반적으로 Unsupevised paradigm이 채택되어 활발히 연구되었지만, 다양한 결함 범주를 구분할 수 없으므로 Classification과 같은 특정 작업에는 적용할 수 없음

  • Data insuffciency problem을 해결하기 위한 직관적인 아이디어는 더 많은 결함 이미지를 생성하는 것이며, 이전의 방법들은 인공물을 수동으로 추가하거나 (DeVries and Taylor 2017) 결함이 없는 이미지의 패치(결함)을 잘라 붙여 넣거나 (Li et al. 2021), 한 이미지에서 다른 이미지로 결함 영역을 복사하는 방식 (Lin et al. 2021)으로 단순하지만 가짜 결함 이미지를 추가하려고 시도했음
    그러나 이러한 방법으로 생성된 결함 이미지는 사실적이고 다양하지 못함

  • 반면, Generatie Adversarial Network (Goodfellow et al. 2014)와 그 variation은 이미지 생성 작업에 널리 사용되고 있지만 data insuffciency problem에 취약하다는 단점때문에 결함 이미지 생성분야에는 몇가지 연구들을 (Niu et al. 2020 Zhang et al. 2021) 제외하고는 거의 사용되지 않음
    해당 연구들은 수백 또는 수천 개의 결함 이미지와 그 이상의 결함 없는 이미지에 의존하거나 단일 Texture 범주에만 초점을 맞추고 있음
    그러나, 실제 산업 제조에서는 생산 라인에서 실제 결함 이미지가 드물고 수집이 어렵기 때문에 일반적으로 소수의 결함 이미지만 활용이 가능함

  • 이러한 한계를 극복하기 위해 본 논문에서는 제안된 결함 이미지를 활용하여 사실적이고 다양한 결함 이미지를 생성하는 새로운 Defect-Aware Feature Manipulation GAN을 제안함
    DFMGAN은 대규모 도메인에서 학습한 Pre-trained 모델을 소규모 도메인에 적용하는 Few-shot image generation method에서 영감을 얻었음
    본 논문에서는 수백 개의 결함 없는 물체 이미지로 훈련된 Backbone generator를 사용한 defect-aware residual blocks을 제안하여 그럴듯한 결함 영역을 생성하고 이 영역 내의 feature를 조작하여 다양하고 사실적인 결함 이미지를 생성함
    Training process는 아래 그림과 같으며, MVTec AD Dataset을 통해 DFMGAN이 다양한 결함 이미지를 고해상도로 생성할 수 있을 뿐만 아니라 비전통적 증강으로 defect inspection 작업의 성능을 향상시킬 수 있음이 입증됨

Method

  • DFMGAN은 두 단계 Training strategy를 채택함
    먼저, 수백개의 결함 없는 이미지에 대해 데이터 효율이 높은 StyleGAN2를 Backbone generator로 훈련하여 결함이 없는 이미지에 임의의 zobjectz_{object}를 매핑함
    그 다음, defect-aware residual blocks를 defect mapping networks와 함께 Backbone generator에 연결하고 추가된 모듈을 몇 개의 결함 이미지에 대해 훈련함
    그 결과, Entire generator는 제어 가능한 결함 영역이 있는 결함 이미지에 zobjectz_{object}와 결함 zdefectz_{defect}를 매핑함

Pretraining on Defect-free Images

  • 첫 번째 훈련 단계에서는 zobjectz_{object}를 무작위로 sampling하여 결함 없는 이미지를 생성할 수 있도록 DFMGAN의 backbone generator로 StyleGAN2를 학습하는 것을 목표로 함
    훈련 결과, backbone generator는 network에서 풍부한 object feature를 encoding하며 추가적인 다음 단계에서 defect-aware residual blocks을 붙여서 결함이 없는 이미지에서 결함이 있는 이미지로 모델을 조정할 수 있음

Transferring to Defect Images

  • 결함 이미지가 결함 영역과 결함이 없는 영역으로 구성된다는 사실을 고려할 때, backbone에서 잠재적 결함 영역을 적절히 조작하면 전체 모델을 확장하여 결함이 없는 이미지의 생성 능력을 유지하면서 결함 이미지를 생성할 수 있을 것임
    이에 영감을 받아 두 번째 훈련 단계에서는 backbone에 부착하는 defect-aware residual blocks을 제안하며 이는 defect mask와 defect feature를 구분하는 그럴듯한 defect mask를 생성함
    그런 다음, mask와 defect feature를 사용하여 object feature map을 조작하여 결함이 없는 이미지에 결함을 추가함
    이 단계에서 결함의 빈도와 변화 등의 다양성을 보장하기 위해 extra defect matching discrimiator와 modified mode seeking loss를 추가로 사용함

Defect-aware Residual Blocks

  • 본 논문에서 제안하는 defect-aware residual blocks는 backbone generator의 synthesis blocks와 유사한 구조를 공유하며, 결함 영역내의 객체 feature만 조작하고 비결함 영역의 feature는 변경되지 않은 상태로 유지시킴
    생성된 결함 이미지의 다양성을 보장하기 위해 결함이 backbone의 object feature map에 의해서만 결정되는 것이 아니라 결함의 변화를 제어하기 위해 eatra defect mapping network를 도입함
    defect mapping network는 무작위로 sampling된 zdefectz_{defect}을 입력받아서 backbone과 유사하게 residual blocks를 변조하는 데 사용되는 가중치 wdefectw_{defect}을 출력함
  • 두 번째 훈련 단계에서 DFMGAN은 backbone generator를 변환하고 결함 이미지에 대해 다양성을 가진 더 많은 defect sample을 생성함
    또한, backbone의 parameter를 고정함으로써 결함 없는 이미지를 생성하는 기능을 유지한다는 점도 DFMGAN의 또 다른 장점임

Two Discriminators

  • 결함이 없는 이미지와 결함이 있는 이미지는 결함 영역을 제외하고는 거의 유사하기 때문에 pre-trained discriminator를 결함이 없는 이미지에서 결함이 있는 이미지로 쉽게 transferring할 수 있음
    그러나, 해당 discriminator는 mask에 대한 판별을 정확하게 할 수 없음. 따라서 생성된 defect mask가 이미지의 결함 영역을 정확하게 구분하도록 보장하기 위해 추가적인 discriminator DmatchD_{match}를 사용하여 간극을 메움
  • DmatchD_{match}는 기존 Discriminator와 거의 동일한 구조를 가지고 있지만, 결함 이미지가 마스크와 일치하는지 여부를 판단하는 것이 더 쉬운 task이므로 각 layer의 channel 수를 줄였음
  • 이미지와 마스크 쌍을 연결한 후, DmatchD_{match}에 입력 시 도출되는 score로 결함 이미지가 해당 결함 마스크와 일치하는지 여부를 판단함

Mode Seeking Loss

  • 지금까지 설명한 DFMGAN의 구조를 통해 예비 실험을 진행한 결과, 생성되는 결함 이미지는 zobjectz_{object}에 많은 영향을 받는다는 것을 알게되었고 이는 유사한 물체에는 항상 닮은 결함이 동반되어 다양성을 크게 해치는 결과를 의미함
    따라서, 이를 완화하기 위해 mode seeking loss를 활용함

  • defect mapping network는 무작위 zdefect1z^{1}_{defect}zdefect2z^{2}_{defect}를 사용하여 두 개의 대응하는 wdefect1w^{1}_{defect}wdefect2w^{2}_{defect}를 출력함
    이를 사용하여 각각의 결함 마스크 M1M_1M2M_2를 생성한 뒤, DFMGAN은 mode seeking loss를 최소화함

  • 즉, 서로 다른 wdefectw_{defect}을 사용할 때 결함 마스크 간의 차이가 최대화되길 바라는 것임

Experiment

Defect Image Generation

Quantitative Result

  • 헤이즐넛의 모든 결함 범주에 대해 DFMGAN은 KID와 Clustered LPIPS 모두에서 다른 모든 방법보다 우수한 성능을 보여, 훈련된 데이터가 매우 부족함에도 불구하고 높은 품질과 다양성을 갖춘 결함 이미지를 생성한다는 것을 보여줌
    참고로 Crop & Paste로 생성된 이미지는 Dataset과 거의 동일한 분포를 가지므로 항상 KID 점수가 거의 0에 가까움

Qualitative Result

  • 위 그림에서 알 수 있듯이 Finetune과 DiffAug로 생성된 이미지의 품질은 좋지만 실제로는 훈련 이미지를 과도하게 왜곡하여 약간의 추가 다양성을 제공하는데, 이와 대조적으로 CDC, SDGAN, Defect-GAN은 특정 물체 범주에서 few-shot learning으로는 사실적인 샘플을 생성할 수 없었음
  • DFMGAN은 현실과 다양성 사이의 균형을 잘 유지하여 만족스러운 이미지를 생성함

Discussion

  • few-shot 이미지 생성 방법(Finetune, DiffAug, CDC)과 이전의 생성 결함 생성 방법(SDGAN, Defect-GAN)을 비교한 결과, 결함 생성을 위한 특별한 설계가 없는 전자는 과적합 경향이 있다는 것을 발견함
    Finetune 또는 DiffAug로 생성된 대부분의 이미지는 실제 결함 이미지 중 하나와 거의 동일하기 때문에 일반적으로 KID score는 좋지만 Clustered LPIPS는 상대적으로 낮음
    반면, CDC는 결함이 없는 이미지가 비슷하게 나타나기 때문에 변화가 많은 결함 이미지를 생성할 수 없다는 한계가 있음

  • 그러나, 앞서 언급한 두 가지 방법에서 가장 중요한 한계점은 이미지 대 이미지 paradigm이라고 저자들은 주장함
    결함을 생성하면서 동시에 결함이 없는 영역의 품질도 보장해야 하는데, 이는 생성된 결함 마스크에 의해 명시적으로 구분된 결함에만 집중하는 것보다 더 어려운 작업임
    따라서, 앞선 두 가지 방법은 더 나쁜 KID score를 도출하며 생성된 이미지가 Downstream task에 도움이 되기에는 현실적이지 못함

  • 또한, Copy&Paste의 경우 극소수의 결함 이미지만을 생성할 수 있으며 때때로 물체 외부에 결함이 존재하는데, 이는 이 방법의 생성 과정이 물체와 외부 요인에 의해 결함이 결정된다는 전제를 위배하기 때문에 적절치 않음

  • 언급한 문제를 해결하기 위해 본 논문의 저자들은 DFMGAN을 설계하였음
    첫 번째 훈련 단계에서는 결함이 없는 이미지로 모델을 훈련하여 결함 범주로 옮기는 데 유리한 물체 범주의 분포를 포착함
    두 번째 훈련 단계에서는 첫 번째 단계에서 학습한 다양한 결함 없는 샘플의 feature에 실제와 같은 결함을 추가하도록 모델을 학습시켜 과적합을 방지하고 생성된 결함 이미지의 다양성을 향상시킴

  • 결론적으로, 특수 설계된 구조를 통해 DFMGAN은 복잡한 구조 정보와 높은 변동성을 가진 물체에서도 까다로운 몇 장의 결함 이미지 생성을 할 수 있으며, 이전 방법보다 성능이 큰 폭으로 향상됨

Data Augmentation for Defect Classification

  • Defect classification은 하나의 물체 범주에서 다양한 유형의 결함을 인식하는 defect inspection임
    Cut&Paste의 한계로 인해 마스크가 필요없는 defect classification에 대해 DFMGAN을 테스트함

  • 위 표는 defect classification의 정확도 결과를 표시한 것인데, 산업현장에서의 실제 defect inspection을 시뮬레이션한 이 classification 작업에서 DFMGAN은 모든 partition에서 가장 높은 점수를 얻었으며, 이는 본 논문의 방법이 Downstream task에서 가장 유익한 data augmentation 기법으로 작용한다는 것을 의미함

Conclusion

  • 본 논문은 소수의 defect image만을 사용해서 다양한 결함 이미지를 고품질로 생성할 수 있는 Few-shot defect image generation 방법인 DFMGAN을 제안함
    DFMGAN은 defect-aware residual blocks이 합리적인 defect mask를 생성하는 방법을 학습하고 그에 따라 물체 feature를 조작하는 것이 특징임
  • MVTec AD Dataset을 활용한 실험을 통해 defect inspection에 대한 이점이 입증됨

0개의 댓글