WaferSegClassNet - A Light-weight Network for Classification and Segmentation of Semiconductor Wafer Defects

This repository contains the source code of our paper, WaferSegClassNet (accepted for publication in Computers in Industry).

we propose WaferSegClassNet (WSCN), a deep-learning model for simultaneously performing both classification and segmentation of defects on Wafer Maps. To the best of our knowledge, WSCN is the first wafer defect analysis model that performs both segmentation and classification. WSCN follows a multi-task learning framework to segment and classify an image simultaneously.

Sample Results

---

Check our project page for more qualitative results.

Project Organization

├── LICENSE
├── README.md          <- The top-level README for developers using this project.
├── data
│   ├── processed      <- The final, canonical data sets for modeling.
│   └── raw            <- The original, immutable data dump.
│
├── requirements.txt   <- The requirements file for reproducing the analysis environment.
│
├── src                <- Source code for use in this project.
│   ├── __init__.py    <- Makes src a Python module
│   ├── config.py      <- All configuration params
|   ├── util.py        <- All utilities functions
│   │
│   ├── data           <- Script to generate data in required format
│   │   └── make_dataset.py
│   │
│   ├── models         <- Scripts to train models and then use trained models to make
│   │   │                 predictions and test performance.
│   │   ├── predict_model.py
│   │   └── train_model.py
|   |   └── test_model.py
|   |   └── network.py
|   |   └── loss.py
├── inference          <- Inference directory, where predicted masks are stored.
├── logs               <- Logs directory for saving terminal output.
├── weights            <- Weights directory for saving checkpoints.

---

Get Started

---

Dependencies:

pip install -r requirements.txt

(Optional) Conda Environment Configuration

First, create a conda environment

conda create -n wscn # python=3
source activate wscn

Now, add dependencies

Now, you can install the required packages.

pip install -r requirements.txt

Dataset

We have used MIXEDWM38 dataset which can be downloaded from here. Download the dataset, unzip it and place Wafer_Map_Datasets.npz in data/raw/Wafer_Map_Datasets.npz.

To prepare the dataset ready for training, Run following command from /src directory.

python data/make_dataset.py

Above command should prepare Images, Labels, and Masks ready for training in data/processed directory.

Training

change the hyperparameters and configuration parameters according to need in src/config.py.

To train wscn, Run following command from /src directory.

python models/train_model.py

All the trained checkpoints for pre-training as well as full model training will be saved in /weights.

Above command will first pre-train encoder with N-Pair contrastive loss and then finetune segmentation and classification for given number of epochs.

Prediction

To train wscn, Run following command from /src directory.

python models/predict_model.py --image <path_of_an_image_in_numpy_format>

Above command will predict the given image and save binary output mask in inference/ directory.

Test performance

To test wscn with trained model, Run following command from /src directory.

python models/test_model.py

Above command will generate IOU Score, and DICE Score for segmentation output, and classification report, Matthews Correlation Coefficient (MCC) and ROC AUC Curve for classification output. ROC-AUC Curve will be saved in inference/ directory.

Citation

@article {ref162,
	title = "WaferSegClassNet - A Light-weight Network for Classification and Segmentation of Semiconductor Wafer Defects",
	year = "2022",
	author = "Subhrajit Nag and Dhruv Makwana and R Sai Chandra Teja and Sparsh Mittal and C Krishna Mohan",
	journal = "Computers in industry",
}

License

---

CC BY-NC-ND 4.0