Open-Set Recognition of Screen Defects with Negative-Guided Augmented Prototype Generator and Open Feature Generation
The official PyTorch implementation of Open-Set Recognition of Screen Defects with Negative-Guided Augmented Prototype Generator and Open Feature Generation.
We have tested in a Python=3.7 environment with PyTorch=1.7.0. Other environments may work as well.
conda create --name FSL python=3.7
conda activate FSL
conda install pytorch==1.7.0 torchvision==0.8.0 cudatoolkit=10.1
pip install terminaltables==3.1.0
pip install scikit-learn==1.0.2
pip install matplotlib==3.5.3
pip install PyYAML==6.0
pip install tqdm==4.64.1
pip install pandas==1.3.4
pip install albumentations==1.0.3
pip install opencv-python==4.5.3.56 #albumentations should be installed before opencv-python
#opencv-python-headless==4.6.0.66
pip install seaborn==0.12.2
Data splits are placed under dataset/
folder.
For the NEU-CLS dataset, you can download from here. Then, you can place images IMAGES/
under dataset/NEU-DET/
folder.
For the PVEL-AD dataset, you can download from here. Then, you can place images ( othertypes/
,trainval/
) under dataset/solar_cell_EL_image/
folder.
There are four parts in the code.
methods
: It contains the open-set recognition method for the experiment, including the entire model, training process, and testing process.models
: It contains the backbone network.data
: Dataloader of different datasets.dataset
: Train/val/test splits for the data setsconfigs
: It contains configuration files for experiments.utils
:
-
Train NEU-CLS
python train_meta.py -default_args_path ./configs/NEU_Ours.yaml # img_path: path to the NEU-CLS dataset. (1) all data splits in img_path/0_OSR_labels/*.csv;(2) images in img_path/IMAGES/*.jpg # split_idx: specify one out of 5 known/unknown class splits
-
Train PVEL-AD
python train_meta.py -default_args_path ./configs/MSD_Ours.yaml # img_path: path to the PVEL-AD dataset. (1) all data splits in img_path/*.csv;(2) images in img_path/othertypes/good/*.jpg, img_path/trainval/JPEGImages/*.jpg
Remember to change img_path
into your own data root, or you will encounter errors.
-
Test NEU-CLS
python test_meta.py -default_args_path ./configs/NEU_Ours.yaml -resume_path ./*.pth # resume_path: the saved weights of the method
-
Test PVEL-AD
python train_meta.py -default_args_path ./configs/MSD_Ours.yaml -resume_path ./*.pth # resume_path: the saved weights of the method
We thank the following repos providing helpful components/functions/dataset in our work.