DOCSEARCHbyAlgolia
Submit Docs

SuperSimpleNet

repository·main·Indexed 41 days ago

https://github.com/blaz-r/supersimplenet

A unified unsupervised and supervised learning framework for fast and reliable surface defect detection. It supports multiple supervision regimes and datasets including MVTec, VisA, KSDD2, and SensumSODF.

Tokens
835
Snippets
4
Records
6
Agent score
2%

What's inside SuperSimpleNet

  1. Prepare datasets for SuperSimpleNet

    main

    Datasets must be prepared and stored in a specific directory structure. You can use the provided scripts in ./datamodules/setup/ to automate the download and extraction for MVTec, VisA, and KSDD2.

    Dataset Preparation Steps:

    1. Navigate to ./datamodules/setup/.
    2. Run prepare_mvtec.py (MVTec).
    3. Run prepare_visa.py (VisA - handles splits automatically).
    4. Run prepare_ksdd2.py (KSDD2).
    5. SensumSODF (Manual):
      • Request a link from the official site.
      • Download and extract to the dataset folder.
      • Download the 3-fold CV split files and place the sensum_splits folder inside the SensumSODF root.

    Required Directory Structure:

    datasets/
    KolektorSDD2/
        train/...
        test/...
    SensumSODF/
        capsule/...
        softgel/...
        sensum_splits/
            capsule/...
            softgel/...
    mvtec/
         bottle/...
    visa/
        visa_pytorch/
            candle/...

  2. Install SuperSimpleNet and set up the environment

    main

    To set up the development environment, create a new Conda environment with Python 3.10 and install the required dependencies via pip.

    Note: The project uses wandb for logging. To enable it, you must uncomment wandb in requirements.txt and set LOG_WANDB=True at the top of train.py.

    conda create -n ssn_env python=3.10
pip install -r requirements.txt

  3. Configure and organize model checkpoints

    main

    Checkpoints can be downloaded from HuggingFace or Google Drive. For the JIMS version, ensure checkpoints are placed in a directory with run_id 0.

    Note on JIMS vs ICPR versions: The JIMS version includes a <ratio> subdirectory (e.g., 246) for mixed supervision scenarios, whereas the original ICPR version does not.

    Expected Checkpoint Path:

    ./weights/
   0/
      ksdd2/
         ksdd2/
            <ratio> (e.g. 246)/
               weights.pt
      sensum/
      mvtec/
      visa/

  5. Run performance benchmarks

    main

    To evaluate performance metrics such as inference speed, throughput, memory consumption, and FLOPs, use the perf_main.py script located in the ./perf directory.

    Pass the GPU model name as an argument.

    For Slurm-based systems, use the run_slurm_perf.sh script.

    python perf_main.py <gpu_model>

  6. Train SuperSimpleNet models

    main

    Train the model by passing the dataset name as a positional argument to train.py.

    Supported dataset names: mvtec, visa, sensum, and ksdd2.

    Configuration:

    • Model and dataset configurations are located within train.py.
    • Parameter Selection:
      • For unsupervised training on your own data: Use MVTec parameters.
      • For supervised training: Use SensumSODF parameters.
    • Key Parameters: dt (distance transform), dilate (label dilation), and adapt_cls_feat (adapted features for cls head).

    For Slurm-based systems, use the run_slurm_train.sh script.

    python train.py <dataset_name>