--- title: Pytorch ViT - SEM Images Defect Detection emoji: 🔍 colorFrom: indigo colorTo: purple sdk: docker app_file: app.py pinned: false --- ## SEM Images Defect Detection using ViT The primary purpose of this project was to learn to apply **ViT**, and experiment any approach not done in classroom. This project involved prompt engineering and was completed with vibe-coding assistance from **Gemini Flash 3**. - **Background**: Semiconductor defect images are scarce & proprietary, and are typically high-res Scanning Electron Microscopy (SEM) images. Similar to images used in this project. Pre-trained models are typically trained with images not as high-res as SEM images. - Workarounds for data scarcity (not applied in this project): Data augmentation or using models to generate synthetic defect images, etc. - **Model training approach**: One-class learning, train on 'normal' data only - [HF app demo](https://huggingface.co/spaces/morinousagi/pytorch-vit-defect-detect) | [UI snapshot](app_ui.png) (Alternate training approach using sample size data - [GitHub repo](https://huggingface.co/spaces/morinousagi/pytorch-vit-defect-detect/blob/main/README.md) ### Data source - Citation: Cheng, Deruo, 2021, "MIIC (Microscopic Images of Integrated Circuits) Dataset for Anomaly Detection and Image Inpainting", https://doi.org/10.21979/N9/WBLTFI, DR-NTU (Data), V1. The data accompanies the publication: "Joint Anomaly Detection and Inpainting for Microscopy Images via Deep Self-Supervised Learning", IEEE International Conference on Image Processing (ICIP), 2021. - File: over 20K normal images and 116 anomaly images -- as it is computationally intensive, only a small subset of this data was used in this project for practice and demo purpose only. ### Dataset - Convert 512x512 grayscale (1 channel) to RGB (3 channels) as expected by the pre-trained ViT - Scaling (handles by ToTensor) Normalization is not required as there is image reconstruction within Autoencoder. Decoder with Sigmoid() will output pixels in [0, 1] -> target to be in the same range as output ``` /data (local) /train /defective (0 images) /normal (3,000 images) /test /defective (116 images) /normal (100 images) ``` ### Model Unsupervised (reconstructive) anomaly detection - Autoencoders - a specialized, self-supervised type (no labels) of encoder-decoder model focused on reconstructing their own input to learn efficient data representations (input = output). - init - Load pre-trained ViT as encoder: `google/vit-base-patch16-224` (expects 224p resolution by default) - Custom convolutional decoder to reconstruct the image - mapping ViT features back to image space - forward pass - Requires position interpolation as input images are maintained at 512p resolution to prevent loss of info (i.e. tiny defects) ### Training & Evaluation - Optimizer: AdamW (improves upon Adam by decoupling weight decay from the gradient update process, for better regularization and generalization) - Learning rate: 1e-4 - Loss: MSE - Evaluation: Peak score logic (1st attempt used average scoring did not yield good results) ### Pipeline Run ```python pipeline.py``` - Set global variables - Run dataloaders, training and evaluation - Calculate best threshold using F1-score - Save model file ### Results Trained with 3000 'normal' images (topographic and metal interconnects) for 15 epochs only - non-concluding ``` --- Starting Training (Normal Images Only) --- Epoch 1: 0%| | 0/375 [00:00