vitjanz/draem

DRAEM

PyTorch implementation of DRAEM - ICCV2021:

@InProceedings{Zavrtanik_2021_ICCV,
    author    = {Zavrtanik, Vitjan and Kristan, Matej and Skocaj, Danijel},
    title     = {DRAEM - A Discriminatively Trained Reconstruction Embedding for Surface Anomaly Detection},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2021},
    pages     = {8330-8339}
}

Datasets

To train on the MVtec Anomaly Detection dataset download the data and extract it. The Describable Textures dataset was used as the anomaly source image set in most of the experiments in the paper. You can run the download_dataset.sh script from the project directory to download the MVTec and the DTD datasets to the datasets folder in the project directory:

./scripts/download_dataset.sh

Training

Pass the folder containing the training dataset to the train_DRAEM.py script as the --data_path argument and the folder locating the anomaly source images as the --anomaly_source_path argument. The training script also requires the batch size (--bs), learning rate (--lr), epochs (--epochs), path to store checkpoints (--checkpoint_path) and path to store logs (--log_path). Example:

python train_DRAEM.py --gpu_id 0 --obj_id -1 --lr 0.0001 --bs 8 --epochs 700 --data_path ./datasets/mvtec/ --anomaly_source_path ./datasets/dtd/images/ --checkpoint_path ./checkpoints/ --log_path ./logs/

The conda environement used in the project is decsribed in requirements.txt.

Pretrained models

Pretrained DRAEM models for each class of the MVTec anomaly detection dataset are available here. To download the pretrained models directly see ./scripts/download_pretrained.sh.

The pretrained models achieve a 98.1 image-level ROC AUC, 97.5 pixel-wise ROC AUC and a 68.9 pixel-wise AP.

Evaluating

The test script requires the --gpu_id arguments, the name of the checkpoint files (--base_model_name) for trained models, the location of the MVTec anomaly detection dataset (--data_path) and the folder where the checkpoint files are located (--checkpoint_path) with pretrained models can be run with:

python test_DRAEM.py --gpu_id 0 --base_model_name "DRAEM_seg_large_ae_large_0.0001_800_bs8" --data_path ./datasets/mvtec/ --checkpoint_path ./checkpoints/DRAEM_checkpoints/

draem's Issues

Positive integer value

data normalization?

Hi, thanks for your great work,
I have a question about data processing,
In most of the implementations in the computer vision area with PyTorch, we will apply image normalization using std and mean,
I wonder why don't you apply normalization in the data preprocessing part?

pre-trained model

I want to train my own dataset, how can I get the file of the pre-trained model?

Pre-trained model on DAGM dataset

Thanks for the excellent work!
Can the authors provide pre-trained models on DAGM? I tried to train on DAGM according to the settings in the paper, but the final Image-level AUC was only 91%.

can not reproduce the results

Hello, I use the checkpoints you offered in the google drive to test the MVTEC datasets , but cannot reproduce the results in the paper. There are large gaps in some classes, but the mean result is the same. The checkpoints you offered are final checkpoints or should be finetuned?Or maybe you made some other changes?
Here are my results

base_model_name of my data

how do create --base_model_name of own data?
how do I predict image in own data?

Making the right environment

Well sorry to interrupt u, but now i have a problem in making the environment. I have tried on Centos 7 and Ubuntu, using the latest Anaconda. There are like 30 or more packages which can not be installed. I have checked online for various kinds of methods, but no one make it through. I wonder maybe u can help me solve it, please.

Error is reported when num_works> 0

TypeError: Caught TypeError in DataLoader worker process 0.
TypeError: 'NoneType' object is not subscriptable

Parallel processing

Hello, I would like to ask you for advice. I want to change the program to two GPUs for parallel processing, but it has been unsuccessful. How can I solve it?

What is image_score?

What is image_score?
I don't understand what that means.

in test_DRAEM.py (line 107)

DRAEM/test_DRAEM.py

Line 107 in 2dbf673

image_score = np.max(out_mask_averaged)

in addition, Can I get F1_score in test(evaluate)?

CUDA out of memory(test_DREAM.py)

Hello.

I tried to run test_DREAM.py, but CUDA out of memory occur. (I use 15.90 GiB memory)

I don't know well about all of your codes, but when i check GPU memory, memory usage continue to increase.

I think some variables need to be allocated at CPU. (by using .cpu().detach())

---

After many attempts, I know where i need to change.

https://github.com/VitjanZ/DRAEM/blob/main/test_DRAEM.py#L96:L99

Before(Original codes)

display_images[cnt_display] = gray_rec[0]
display_gt_images[cnt_display] = gray_batch[0]
display_out_masks[cnt_display] = t_mask[0]
display_in_masks[cnt_display] = true_mask[0]

After(changed codes)

display_images[cnt_display] = gray_rec[0].cpu().detach()
display_gt_images[cnt_display] = gray_batch[0].cpu().detach()
display_out_masks[cnt_display] = t_mask[0].cpu().detach()
display_in_masks[cnt_display] = true_mask[0].cpu().detach()

Inference

How to determine the threshold of the discriminant network in the inference process

How to make a heatmap?

Hello
How do I generate a heatmap in test_DRAEM.py when I predict data?

the error in the 124th row of data_loader

The generated image in Figure 4 of the original paper does not correspond to the code in the data loader. The correct code should delete /255 in the 124th row of data_loader. I hope that the author corrects it to avoid more misleading.

Data indicators

May I ask if the indicators imageAUROC, pixel AUROC, and pixel AP in your paper were tested with the same or different weights?

Docker image

Hi, could you possibly provide a docker file for easy library installation? The requirements.txt is causing countless dependency conflicts. Thank you.

Inference

how to inference images with the code?
Thank you

Question about the requirement.txt

The format of requirement is strange and can't be used as "pip install -r" or "conda create -n -f", may be a yaml file is more suitable?
I'll appreciate it if the available environment information provided.
Thanks a lot.