🤖 UniNet: A Contrastive Learning-guided Unified Framework with Feature Selection for Anomaly Detection

Shun Wei,  Jielin Jiang,  Xiaolong Xu

-- CVPR 2025 --

          📄 PDF

🔔 News

• 05-15-2025: Fix some minor errors.
• 04-27-2025: The checkpoints have been uploaded.
• 04-26-2025: The code has been released 😊!
• 04-21-2025: The code will be released in recent days !
• 02-27-2025: Accepted by CVPR2025 💥!

List of TODOs

• 📖 Introduction
• 🔧 Environments
• 📊 Data Preparation
• 🚀 Run Experiments
• 📂 Results
• 🔗 Citation
• 🙏 Acknowledgements
• 📜 License
• 💬 Contact

📖 Introduction

This repository will contain source code for UniNet implemented with PyTorch.

UniNet is a unified framework designed for diverse domains, such as industrial, medical, and video, by addressing the limitations in general ability of existing methods. Besides, UniNet can be suited for both unsupervised and supervised settings simultaneously. As a unified framework, UniNet also focuses on multi-class anomaly detection.

🔧 Environments

Create a new conda environment and install required packages.

conda create -n uninet_env python=3.9.7
conda activate uninet_env
pip install -r requirements.txt

Experiments are conducted on NVIDIA GeForce RTX 3090. Same GPU and package version are recommended.

📊 Data Preparation

The public datasets employed in the paper are listed below.

• Industrial domain: MVTec AD, VisA, BTAD, MVTec 3D-AD, VAD.
• Medical domain: OCT2017, APTOS, ISIC2018, Kvasir, CVC-ColonDB, CVC-ClinicDB.
• Video domain: Ped2.

1. MVTec AD

|-- mvtec
    |-- bottle
        |-- ground_truth
        |-- test
        |-- train
    |-- cable
        |-- ground_truth
        |-- test
        |-- train
    |-- ...

2. VisA

Unzip the file to ../VisA/. Preprocess the dataset to ../VisA_pytorch/ in 1-class mode by their official splitting code. Alternately, you can also preprocess the dataset using this code from ReContrast.

|-- VisA_pytorch
    |-- 1cls
        |-- candle
            |-- ground_truth
            |-- test
                    |-- good
                    |-- bad
            |-- train
                    |-- good
        |-- capsules
        |-- ....

3. APTOS

Creat a new directory ../APTOS. Unzip the file to ../APTOS/original/. Now, the directory would be like:

|-- APTOS
    |-- original
        |-- test_images
        |-- train_images
        |-- test.csv
        |-- train.csv

Run the following command to preprocess the data to ../APTOS/.

python ./prepare_data/prepare_aptos.py --data-folder ../APTOS/original --save-folder ../APTOS

Then, it woule be like:

|-- APTOS
    |-- test
        |-- NORMAL
        |-- ABNORMAL
    |-- train
        |-- NORMAL
    |-- original

4. OCT2017

Creat a new directory ../OCT2017. Unzip the file, and move everything in ZhangLabData/CellData/OCT to ../OCT2017/. The directory should be like:

|-- OCT2017
    |-- test
        |-- CNV
        |-- DME
        |-- DRUSEN
        |-- NORMAL
    |-- train
        |-- CNV
        |-- DME
        |-- DRUSEN
        |-- NORMAL

5. ISIC2018

Creat a new directory ../ISIC2018. After downloading "Training Data","Training Ground Truth", "Validation Data", and "Validation Ground Truth" of Task 3, please unzip them to ../ISIC2018/original/. Now, the directory would be like:

|-- ISIC2018
    |-- original
        |-- ISIC2018_Task3_Training_GroundTruth
        |-- ISIC2018_Task3_Training_Input
        |-- ISIC2018_Task3_Validation_GroundTruth
        |-- ISIC2018_Task3_Validation_Input

Run the following command to preprocess the data to ../ISIC2018/.

python ./prepare_data/prepare_isic2018.py --data-folder ../ISIC2018/original --save-folder ../ISIC2018

Then, it would be like:

|-- ISIC2018
    |-- test
        |-- NORMAL
        |-- ABNORMAL
    |-- train
        |-- NORMAL
    |-- original

🚀 Run Experiments

1. Unsupervised AD

Run the following command for industrial domain, such as MVTec AD dataset:

python main.py --setting oc --dataset "MVTec AD"

Run the following command for medical domain, such as APTOS dataset:

python main.py --dataset APTOS

Alternatively, you can train and test all medical datasets at once using the following command:

python main.py --dataset APTOS --train_and_test_all

Run the following command for video domain:

python main.py --dataset Ped2

Run the following command for multiclass anomaly detection, such as VisA dataset:

python main.py --setting mc --dataset VisA

2. Supervised AD

Run the following command for medical polyp segmentation, sucha as Kvasir-SEG dataset:

python main.py --dataset "Kvasir-SEG"

Similarly, you can also train and test all medical datasets at once using the following command:

python main.py --dataset "Kvasir-SEG" --train_and_test_all

Run the following command for industrial domain:

python main.py --dataset VAD

3. Testing

With training on your own or saving the weight file we upload, you can test the model using the following command:

python main.py --setting oc --dataset "MVTec AD" --load_ckpts

python main.py --setting mc --dataset VisA --load_ckpts

python main.py --dataset APTOS --load_ckpts

For unsupervised industrial AD, the ''--setting'' must be considered.

📂 Results

1. Unsupervised anomaly detection

2. Supervised anomaly detection

3. Multi-class setting

4. Visualization results

🔗 Citation

Please cite our paper if the method and code help and inspire your project:

@inproceedings{wei2025uninet,
  title={UniNet: A Contrastive Learning-guided Unified Framework with Feature Selection for Anomaly Detection},
  author={Wei, Shun and Jiang, Jielin and Xu, Xiaolong},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2025},
}

🙏 Acknowledgements

We sincerely appreciate ReContrast for its clear and excellent approach.

📜 License

The code in this repository is licensed under the MIT license.

💬 Contact

For the questions or collaborations:

• Email: [email protected]
• GitHub issues: open issue

⭐Please consider giving me a star if you find it useful!!!