CAM-UNET: Class Activation MAP Guided UNET with Feedback Refinement for Defect Segmentation
Dongyun Lin, Yiqun Li, Shitala Prasad, Tin Lay Nwe, Sheng Dong, Zaw Min Oo
Abstract
This paper tackles the task of defect segmentation by exploiting sufficient normal (defect-free) training images and limited annotated anomalous images. We propose a class activation map guided UNet (CAM-UNet) with feedback refinement mechanism for accurate defect segmentation. We first modify and pretrain the encoder of a VGG-16 backboned UNet to classify normal and anomalous training images. Then, for each of the anomalous training images, a CAM is generated as the prior segmentation information. Based on the CAM, we propose a feedback refinement process to train two decoder networks to progressively improve the segmentation output. Extensive experiments conducted on MVTEC AD dataset show that the proposed method significantly outperforms multiple benchmarking UNet methods in terms of mean IOU.