/OSR_Screen

Primary LanguagePythonMIT LicenseMIT

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.

Prerequisites

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

Dataset

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.

Code Structures

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 sets
  • configs: It contains configuration files for experiments.
  • utils:

Training scripts

  • 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 scripts

  • 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

Acknowledgment

We thank the following repos providing helpful components/functions/dataset in our work.