shijianjian/dia

Dissolving Is Amplifying: Towards Fine-Grained Anomaly Detection https://arxiv.org/pdf/2302.14696

Dissolving Is Amplifying: Towards Fine-Grained Anomaly Detection

Jian Shi¹ · Pengyi Zhang² · Ni Zhang² · Hakim Ghazzai¹ · Peter Wonka¹

¹King Abdullah University of Science and Technology
²NEC Laboratories China

DIA is a fine-grained anomaly detection framework for medical images. We describe two novel components.

• First, the dissolving transformations. Our main observation is that generative diffusion models are feature-aware and applying them to medical images in a certain manner can remove or diminish fine-grained discriminative features such as tumors or hemorrhaging. More visual demonstration about the dissolving effects are here.
• Second, an amplifying framework. It is based on contrastive learning to learn a semantically meaningful representation of medical images in a self-supervised manner.

The amplifying framework contrasts additional pairs of images with and without dissolving transformations applied and thereby boosts the learning of fine-grained feature representations. DIA significantly improves the medical anomaly detection performance with around 18.40% AUC boost against the baseline method and achieves an overall SOTA against other benchmark methods.

1. Requirements

Environments

$ pip install -r requirements.txt

Datasets

We majorly use the following datasets to benchmark our method:

• APTOS2019,
• SARS-COV-2
• Retina OCT
• Kvasir

2. Training

Pretraining of diffusion models

Modify the data paths in diffusion_models/train_eval_ddpm.py, then run:

python train_eval_ddpm.py --mode train --dataset <DATASET>

In our experiments, the diffusion models are usable within 5 epochs.

Training DIA

To train unlabeled one-class & multi-class models in the paper, run this command:

python train.py --data_root ../data --dataset <DATASET> --model resnet18_imagenet --mode simclr_DIA --shift_trans_type diffusion_rotation --diff_resolution 32 --batch_size 32 --one_class_idx 0 --save_step 1  --diffusion_model_path <DIFFUSION_MODEL_PATH>

• Other transformation types are available to use, e.g. diffusion_rotation, diffusion_cutperm, blurgaussian_rotation, blurgaussian_cutperm, blurmedian_rotation, blurmedian_cutperm. Note that cutperm may perform better if your dataset is more rotation-invariant, while the rotation may perform better if the dataset is more "permutation-invariant".
• To test different resolutions, do remember to re-train the diffusion models for the corresponding resolution.
• For low resolution datasets, e.g. CIFAR10, use --model resnet18 instead of --model resnet18_imagenet.

3. Evaluation

During the evaluation, we use only rotation shift transformation for evaluation.

python eval.py --mode ood_pre --dataset <DATASET> --model resnet18_imagenet --ood_score CSI --shift_trans_type rotation --print_score --diffusion_upper_offset 0. --diffusion_offset 0. --ood_samples 10 --resize_factor 0.54 --resize_fix --one_class_idx 0 --load_path <MODEL_PATH>

4. Results

In general, our method achieves SOTA in terms of fine-grained anomaly detection tasks. We recommend using 32x32 diffusion_rotation for most tasks, as the following results:

Different transformations

In general, diffusion_rotation performs best.

Dataset	transform	Gaussian	Median	Diffusion
SARS-COV-2	Perm	0.788	0.826	0.841
	Rotate	0.847	0.837	0.851
Kvasir	Perm	0.712	0.663	0.840
	Rotate	0.739	0.687	0.860
Retinal OCT	Perm	0.754	0.747	0.890
	Rotate	0.895	0.876	0.919
APTOS 2019	Perm	0.942	0.929	0.926
	Rotate	0.922	0.918	0.934

Resolutions

Higher feature dissolving resolution will dramatically increase the processing time, while hardly bringing up the detection performances.

Dataset	resolution	MACs (G)	AUC
SARS-COV-2	32	2.33	0.851
	64	3.84	0.803
	128	9.90	0.807
Kvasir	32	2.33	0.860
	64	3.84	0.721
	128	9.90	0.730
Retinal OCT	32	2.33	0.919
	64	3.84	0.922
	128	9.90	0.930
APTOS 2019	32	2.33	0.934
	64	3.84	0.937
	128	9.90	0.905

Citation

@misc{2302.14696,
  Author = {Jian Shi and Pengyi Zhang and Ni Zhang and Hakim Ghazzai and Peter Wonka},
  Title = {Dissolving Is Amplifying: Towards Fine-Grained Anomaly Detection},
  Year = {2023},
  Eprint = {arXiv:2302.14696},
}

Acknowledgement

Our method is based on CSI.

dia's Issues

Problems with paper

Hi,

1. What is the difference between fine-grained images and anomaly images? Can you explain what is finegrained parts?
2. Why do you use diffusion? Why didn't you use GAN? Or even autoencoder?
3. You said we removed the fine-grained part of the image with diffusion; how does it help the contrast learning detect whether the image is an anomaly?
4. Our diffusion is trained on the healthy images, and we remove the fine-grained part in DIA. I can't get this part. I mean the difference between the anomaly part and the fine-grained part.
   Because our diffusion is trained on helathy images so we use healthy images for training of the DIA. how it can remove fine-grained parts?

stable diffusion

Hi Jian,
Have you ever tested your model with stable diffusion?

error

I have this error:
File "/content/drive/MyDrive/BS.c_project/code/train.py", line 6, in
from common.train import *
File "/content/drive/MyDrive/BS.c_project/code/common/train.py", line 10, in
import models.classifier as C
File "/content/drive/MyDrive/BS.c_project/code/models/classifier.py", line 5, in
from models.resnet_imagenet_simsiam import resnet18_simsiam, resnet50_simsiam
ModuleNotFoundError: No module named 'models.resnet_imagenet_simsiam'

training

Hi,
I trained your model with pneumoniamnist, but the result in the eval section does not match with your result in the paper(my result is 0.80)
can you tell me what hyperparameters did you use or what should I do?

Readme

Hi,
you miss --data_root in the readme for the eval part command.

resume_path

Another problem is in the resume path which has an error :

=> Loading checkpoint from '/content/drive/MyDrive/BS.c_project/DIA/logs/pneumoniamnist_resnet18_unsup_simclr_DIA_shift_DIA_diffusion_rotation_res32_one_class_0'
Traceback (most recent call last):
File "/content/drive/MyDrive/BS.c_project/DIA/train.py", line 6, in
from common.train import *
File "/content/drive/MyDrive/BS.c_project/DIA/common/train.py", line 132, in
best = config['best']
KeyError: 'best'

saved model

Hi Jian,
I have a question for you, the model of diffusion you saved has keys that we use model key for training. In other codes when we want to save the model we have these keys

odict_keys(['time_embed.0.weight', 'time_embed.0.bias', 'time_embed.2.weight', 'time_embed.2.bias', 'input_blocks.0.0.weight', 'input_blocks.0.0.bias', 'input_blocks.1.0.in_layers.0.weight', 'input_blocks.1.0.in_layers.0.bias', 'input_blocks.1.0.in_layers.2.weight', 'input_blocks.1.0.in_layers.2.bias', 'input_blocks.1.0.emb_layers.1.weight', 'input_blocks.1.0.emb_layers.1.bias', ........

In the other hand your model keys are
"betas", "alphas_cumprod", "alphas_cumprod_prev", "sqrt_alphas_cumprod"

So I get this error while trying other saved models:
RuntimeError: Error(s) in loading state_dict for GaussianDiffusion:
Missing key(s) in state_dict: "betas", "alphas_cumprod", "alphas_cumprod_prev", "sqrt_alphas_cumprod", "sqrt_one_minus_alphas_cumprod", "log_one_minus_alphas_cumprod", "sqrt_recip_alphas_cumprod", ......
Unexpected key(s) in state_dict: "time_embed.0.weight", "time_embed.0.bias", "time_embed.2.weight", "time_embed.2.bias", ."input_blocks.5.0.in_layers.0.weight", "input_blocks.5.0.in_layers.0.bias", "input_blocks.5.0.in_layers.2.weight", ........

So how can I save my model similar to your pretrained diffusion keys?

datasets

Hi,
What does the dataset directory should look like?