first commit

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.txt.gz filter=lfs diff=lfs merge=lfs -text
+weights/pretrained_all.pth filter=lfs diff=lfs merge=lfs -text
+weights/pretrained_mvtec_colondb.pth filter=lfs diff=lfs merge=lfs -text
+weights/pretrained_visa_clinicdb.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+/result/
+/.idea/
+/__pycache__/
+/weights/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Yunkang Cao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,176 @@
----
-title: AdaCLIP
-emoji: 🌖
-colorFrom: pink
-colorTo: pink
-sdk: gradio
-sdk_version: 4.38.1
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# AdaCLIP (Detecting Anomalies for Novel Categories)
+[![HuggingFace Space](https://img.shields.io/badge/🤗-HuggingFace%20Space-cyan.svg)]()
+
+> [**ECCV 24**] [**AdaCLIP: Adapting CLIP with Hybrid Learnable Prompts for Zero-Shot Anomaly Detection**]().
+>
+> by [Yunkang Cao](https://caoyunkang.github.io/), [Jiangning Zhang](https://zhangzjn.github.io/),  [Luca Frittoli](https://scholar.google.com/citations?user=cdML_XUAAAAJ), 
+> [Yuqi Cheng](https://scholar.google.com/citations?user=02BC-WgAAAAJ&hl=en), [Weiming Shen](https://scholar.google.com/citations?user=FuSHsx4AAAAJ&hl=en), [Giacomo Boracchi](https://boracchi.faculty.polimi.it/) 
+> 
+
+## Introduction 
+Zero-shot anomaly detection (ZSAD) targets the identification of anomalies within images from arbitrary novel categories. 
+This study introduces AdaCLIP for the ZSAD task, leveraging a pre-trained vision-language model (VLM), CLIP. 
+AdaCLIP incorporates learnable prompts into CLIP and optimizes them through training on auxiliary annotated anomaly detection data. 
+Two types of learnable prompts are proposed: \textit{static} and \textit{dynamic}. Static prompts are shared across all images, serving to preliminarily adapt CLIP for ZSAD. 
+In contrast, dynamic prompts are generated for each test image, providing CLIP with dynamic adaptation capabilities. 
+The combination of static and dynamic prompts is referred to as hybrid prompts, and yields enhanced ZSAD performance. 
+Extensive experiments conducted across 14 real-world anomaly detection datasets from industrial and medical domains indicate that AdaCLIP outperforms other ZSAD methods and can generalize better to different categories and even domains. 
+Finally, our analysis highlights the importance of diverse auxiliary data and optimized prompts for enhanced generalization capacity.
+
+## Overview of AdaCLIP
+![overview](asset/framework.png)
+
+## 🛠️ Getting Started
+
+### Installation
+To set up the AdaCLIP environment, follow one of the methods below:
+
+- Clone this repo:
+  ```shell
+  git clone https://github.com/caoyunkang/AdaCLIP.git && cd AdaCLIP
+  ```
+- You can use our provided installation script for an automated setup::
+  ```shell
+  sh install.sh
+  ```
+- If you prefer to construct the experimental environment manually, follow these steps:
+  ```shell
+  conda create -n AdaCLIP python=3.9.5 -y
+  conda activate AdaCLIP
+  pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu111/torch_stable.html
+  pip install tqdm tensorboard setuptools==58.0.4 opencv-python scikit-image scikit-learn matplotlib seaborn ftfy regex numpy==1.26.4
+  pip install gradio # Optional, for app 
+  ```
+- Remember to update the dataset root in config.py according to your preference:
+  ```python
+  DATA_ROOT = '../datasets' # Original setting
+  ```
+
+### Dataset Preparation 
+Please download our processed visual anomaly detection datasets to your `DATA_ROOT` as needed. 
+
+#### Industrial Visual Anomaly Detection Datasets
+Note: some links are still in processing...
+
+| Dataset | Google Drive | Baidu Drive | Task
+|------------|------------------|------------------| ------------------|
+| MVTec AD    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization |
+| VisA    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization |
+| MPDD    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization | 
+| BTAD    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization |
+| KSDD    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization |
+| DAGM    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization |
+| DTD-Synthetic    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection & Localization |
+
+
+
+
+#### Medical Visual Anomaly Detection Datasets
+| Dataset | Google Drive | Baidu Drive | Task
+|------------|------------------|------------------|  ------------------|
+| HeadCT    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection |
+| BrainMRI    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection |
+| Br35H    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Detection |
+| ISIC    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Localization |
+| ColonDB    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Localization |
+| ClinicDB    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Localization |
+| TN3K    | [Google Drive](链接) | [Baidu Drive](链接) | Anomaly Localization |
+
+#### Custom Datasets
+To use your custom dataset, follow these steps:
+
+1. Refer to the instructions in `./data_preprocess` to generate the JSON file for your dataset.
+2. Use `./dataset/base_dataset.py` to construct your own dataset.
+
+
+### Weight Preparation
+
+We offer various pre-trained weights on different auxiliary datasets. 
+Please download the pre-trained weights in `./weights`.
+
+| Pre-trained Datasets | Google Drive | Baidu Drive 
+|------------|------------------|------------------|  
+| MVTec AD & ClinicDB    | [Google Drive](https://drive.google.com/file/d/1xVXANHGuJBRx59rqPRir7iqbkYzq45W0/view?usp=drive_link) | [Baidu Drive](链接) | 
+| VisA & ColonDB    | [Google Drive](https://drive.google.com/file/d/1QGmPB0ByPZQ7FucvGODMSz7r5Ke5wx9W/view?usp=drive_link) | [Baidu Drive](链接) | 
+| All Datasets Mentioned Above   | [Google Drive](https://drive.google.com/file/d/1Cgkfx3GAaSYnXPLolx-P7pFqYV0IVzZF/view?usp=drive_link) | [Baidu Drive](链接) |
+
+
+### Train
+
+By default, we use MVTec AD & ClinicDB for training and VisA for validation:
+```shell
+CUDA_VISIBLE_DEVICES=0 python train.py --save_fig True --training_data mvtec colondb --testing_data visa
+```
+
+
+Alternatively, for evaluation on MVTec AD & ClinicDB, we use VisA & ColonDB for training and MVTec AD for validation.
+```shell
+CUDA_VISIBLE_DEVICES=0 python train.py --save_fig True --training_data visa clinicdb --testing_data mvtec
+```
+Since we have utilized half-precision (FP16) for training, the training process can occasionally be unstable.
+It is recommended to run the training process multiple times and choose the best model based on performance
+on the validation set as the final model.
+
+
+To construct a robust ZSAD model for demonstration, we also train our AdaCLIP on all AD datasets mentioned above:
+```shell
+CUDA_VISIBLE_DEVICES=0 python train.py --save_fig True \
+--training_data \
+br35h brain_mri btad clinicdb colondb \
+dagm dtd headct isic mpdd mvtec sdd tn3k visa \
+--testing_data mvtec
+```
+
+### Test
+
+Manually select the best models from the validation set and place them in the `weights/` directory. Then, run the following testing script:
+```shell
+sh test.sh
+```
+
+If you want to test on a single image, you can refer to `test_single_image.sh`:
+```shell
+CUDA_VISIBLE_DEVICES=0 python test.py --testing_model image --ckt_path weights/pretrained_all.pth --save_fig True \
+ --image_path asset/img.png --class_name candle --save_name test.png
+```
+
+## Main Results
+
+Due to differences in versions utilized, the reported performance may vary slightly compared to the detection performance 
+with the provided pre-trained weights. Some categories may show higher performance while others may show lower.
+
+![Table_industrial](./asset/Table_industrial.png)
+![Table_medical](./asset/Table_medical.png)
+![Fig_detection_results](./asset/Fig_detection_results.png)
+
+### :page_facing_up: Demo App
+
+To run the demo application, use the following command:
+
+```bash
+python app.py
+```
+
+![Demo](./asset/Fig_app.png)
+
+## 💘 Acknowledgements
+Our work is largely inspired by the following projects. Thanks for their admiring contribution.
+
+- [VAND-APRIL-GAN](https://github.com/ByChelsea/VAND-APRIL-GAN)
+- [AnomalyCLIP](https://github.com/zqhang/AnomalyCLIP)
+- [SAA](https://github.com/caoyunkang/Segment-Any-Anomaly)
+
+
+## Stargazers over time
+[![Stargazers over time](https://starchart.cc/caoyunkang/AdaCLIP.svg?variant=adaptive)](https://starchart.cc/caoyunkang/AdaCLIP)
+
+
+## Citation
+
+If you find this project helpful for your research, please consider citing the following BibTeX entry.
+
+```BibTex
+
+
+
+```

app.py

@@ -0,0 +1,133 @@
+import gradio as gr
+from PIL import Image, ImageDraw, ImageFont
+import warnings
+import os
+os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+import json
+import os
+import torch
+from scipy.ndimage import gaussian_filter
+import cv2
+from method import AdaCLIP_Trainer
+import numpy as np
+
+############ Init Model
+ckt_path1 = 'weights/pretrained_mvtec_colondb.pth'
+ckt_path2 = "weights/pretrained_visa_clinicdb.pth"
+ckt_path3 = 'weights/pretrained_all.pth'
+
+# Configurations
+image_size = 518
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# device = 'cpu'
+model = "ViT-L-14-336"
+prompting_depth = 4
+prompting_length = 5
+prompting_type = 'SD'
+prompting_branch = 'VL'
+use_hsf = True
+k_clusters = 20
+
+config_path = os.path.join('./model_configs', f'{model}.json')
+
+# Prepare model
+with open(config_path, 'r') as f:
+    model_configs = json.load(f)
+
+# Set up the feature hierarchy
+n_layers = model_configs['vision_cfg']['layers']
+substage = n_layers // 4
+features_list = [substage, substage * 2, substage * 3, substage * 4]
+
+model = AdaCLIP_Trainer(
+    backbone=model,
+    feat_list=features_list,
+    input_dim=model_configs['vision_cfg']['width'],
+    output_dim=model_configs['embed_dim'],
+    learning_rate=0.,
+    device=device,
+    image_size=image_size,
+    prompting_depth=prompting_depth,
+    prompting_length=prompting_length,
+    prompting_branch=prompting_branch,
+    prompting_type=prompting_type,
+    use_hsf=use_hsf,
+    k_clusters=k_clusters
+).to(device)
+
+
+def process_image(image, text, options):
+    # Load the model based on selected options
+    if 'MVTec AD+Colondb' in options:
+        model.load(ckt_path1)
+    elif 'VisA+Clinicdb' in options:
+        model.load(ckt_path2)
+    elif 'All' in options:
+        model.load(ckt_path3)
+    else:
+        # Default to 'All' if no valid option is provided
+        model.load(ckt_path3)
+        print('Invalid option. Defaulting to All.')
+
+    # Ensure image is in RGB mode
+    image = image.convert('RGB')
+
+    # Convert PIL image to NumPy array
+    np_image = np.array(image)
+
+    # Convert RGB to BGR for OpenCV
+    np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
+    np_image = cv2.resize(np_image, (image_size, image_size))
+    # Preprocess the image and run the model
+    img_input = model.preprocess(image).unsqueeze(0)
+    img_input = img_input.to(model.device)
+
+    with torch.no_grad():
+        anomaly_map, anomaly_score = model.clip_model(img_input, [text], aggregation=True)
+
+    # Process anomaly map
+    anomaly_map = anomaly_map[0, :, :].cpu().numpy()
+    anomaly_score = anomaly_score[0].cpu().numpy()
+    anomaly_map = gaussian_filter(anomaly_map, sigma=4)
+    anomaly_map = (anomaly_map * 255).astype(np.uint8)
+
+    # Apply color map and blend with original image
+    heat_map = cv2.applyColorMap(anomaly_map, cv2.COLORMAP_JET)
+    vis_map = cv2.addWeighted(heat_map, 0.5, np_image, 0.5, 0)
+
+    # Convert OpenCV image back to PIL image for Gradio
+    vis_map_pil = Image.fromarray(cv2.cvtColor(vis_map, cv2.COLOR_BGR2RGB))
+
+    return vis_map_pil, f'{anomaly_score:.3f}'
+
+# Define examples
+examples = [
+    ["asset/img.png", "candle", "MVTec AD+Colondb"],
+    ["asset/img2.png", "bottle", "VisA+Clinicdb"],
+    ["asset/img3.png", "button", "All"],
+]
+
+# Gradio interface layout
+demo = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Class Name"),
+        gr.Radio(["MVTec AD+Colondb",
+                  "VisA+Clinicdb",
+                  "All"],
+        label="Pre-trained Datasets")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Output Image"),
+        gr.Textbox(label="Anomaly Score"),
+    ],
+    examples=examples,
+    title="AdaCLIP -- Zero-shot Anomaly Detection",
+    description="Upload an image, enter class name, and select pre-trained datasets to do zero-shot anomaly detection"
+)
+
+# Launch the demo
+demo.launch()
+# demo.launch(server_name="0.0.0.0", server_port=10002)
+

config.py

@@ -0,0 +1 @@
+DATA_ROOT = '../datasets'

data_preprocess/br35h.py

@@ -0,0 +1,50 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+Br35h_ROOT = os.path.join(DATA_ROOT, 'Br35h_anomaly_detection')
+class Br35hSolver(object):
+    CLSNAMES = [
+        'br35h',
+    ]
+
+    def __init__(self, root=Br35h_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    img_names.sort()
+
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = Br35hSolver(root=Br35h_ROOT)
+    runner.run()

data_preprocess/brain_mri.py

@@ -0,0 +1,51 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+BrainMRI_ROOT = os.path.join(DATA_ROOT, 'BrainMRI')
+
+class BrainMRISolver(object):
+    CLSNAMES = [
+        'brain_mri',
+    ]
+
+    def __init__(self, root=BrainMRI_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    img_names.sort()
+
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = BrainMRISolver(root=BrainMRI_ROOT)
+    runner.run()

data_preprocess/btad.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+BTAD_ROOT = os.path.join(DATA_ROOT, 'BTech_Dataset_transformed')
+
+class BTADSolver(object):
+    CLSNAMES = [
+        '01', '02', '03',
+    ]
+
+    def __init__(self, root=BTAD_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['ok'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = BTADSolver(root=BTAD_ROOT)
+    runner.run()

data_preprocess/clinicdb.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+ClinicDB_ROOT = os.path.join(DATA_ROOT, 'CVC-ClinicDB')
+
+class ClinicDBSolver(object):
+    CLSNAMES = [
+        'ClinicDB',
+    ]
+
+    def __init__(self, root=ClinicDB_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = ClinicDBSolver(root=ClinicDB_ROOT)
+    runner.run()

data_preprocess/colondb.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+ColonDB_ROOT = os.path.join(DATA_ROOT, 'CVC-ColonDB')
+
+class ColonDBSolver(object):
+    CLSNAMES = [
+        'ColonDB',
+    ]
+
+    def __init__(self, root=ColonDB_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = ColonDBSolver(root=ColonDB_ROOT)
+    runner.run()

data_preprocess/dagm-pre.py

@@ -0,0 +1,82 @@
+import os
+import numpy as np
+from sklearn.model_selection import train_test_split
+import cv2
+import argparse
+from config import DATA_ROOT
+
+dataset_root = os.path.join(DATA_ROOT, 'DAGM2007')
+
+class_names = os.listdir(dataset_root)
+
+
+for class_name in class_names:
+    states = os.listdir(os.path.join(dataset_root, class_name))
+    for state in states:
+        images = list()
+        mask = list()
+        files = os.listdir(os.path.join(dataset_root, class_name,state))
+        for f in files:
+            if 'PNG' in f[-3:]:
+                images.append(f)
+        files = os.listdir(os.path.join(dataset_root, class_name, state,'Label'))
+        for f in files:
+            if 'PNG' in f[-3:]:
+                mask.append(f)
+        normal_image_path_train = list()
+        normal_image_path_test = list()
+        normal_image_path = list()
+        abnormal_image_path = list()
+        abnormal_image_label = list()
+        for f in images:
+            id = f[-8:-4]
+            flag = 0
+            for y in mask:
+                if id in y:
+                    abnormal_image_path.append(f)
+                    abnormal_image_label.append(y)
+                    flag = 1
+                    break
+            if flag == 0:
+                normal_image_path.append(f)
+
+        if len(abnormal_image_path) != len(abnormal_image_label):
+            raise ValueError
+        length = len(abnormal_image_path)
+
+        normal_image_path_test = normal_image_path[:length]
+        normal_image_path_train = normal_image_path[length:]
+
+        target_root = '../datasets/DAGM_anomaly_detection'
+
+        train_root = os.path.join(target_root, class_name, 'train','good')
+        if not os.path.exists(train_root):
+            os.makedirs(train_root)
+        for f in normal_image_path_train:
+            image_data = cv2.imread(os.path.join(dataset_root, class_name, state,f))
+            cv2.imwrite(os.path.join(train_root,f), image_data)
+
+        test_root = os.path.join(target_root, class_name, 'test','good')
+        if not os.path.exists(test_root):
+            os.makedirs(test_root)
+        for f in normal_image_path_test:
+            image_data = cv2.imread(os.path.join(dataset_root, class_name, state,f))
+            cv2.imwrite(os.path.join(test_root,f), image_data)
+
+        test_root = os.path.join(target_root, class_name, 'test','defect')
+        if not os.path.exists(test_root):
+            os.makedirs(test_root)
+        for f in abnormal_image_path:
+            image_data = cv2.imread(os.path.join(dataset_root, class_name, state,f))
+            cv2.imwrite(os.path.join(test_root,f), image_data)
+
+        test_root = os.path.join(target_root, class_name, 'ground_truth','defect')
+        if not os.path.exists(test_root):
+            os.makedirs(test_root)
+        for f in mask:
+            image_data = cv2.imread(os.path.join(dataset_root, class_name, state,'Label',f))
+            cv2.imwrite(os.path.join(test_root,f), image_data)
+
+
+
+print("Done")

data_preprocess/dagm.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+DAGM_ROOT = os.path.join(DATA_ROOT, 'DAGM_anomaly_detection')
+
+class DAGMSolver(object):
+    CLSNAMES = [
+        'Class1', 'Class2', 'Class3', 'Class4', 'Class5','Class6','Class7','Class8','Class9','Class10',
+    ]
+
+    def __init__(self, root=DAGM_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = DAGMSolver(root=DAGM_ROOT)
+    runner.run()

data_preprocess/dtd.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+DTD_ROOT = os.path.join(DATA_ROOT, 'DTD-Synthetic')
+
+class DTDSolver(object):
+    CLSNAMES = [
+        'Blotchy_099', 'Fibrous_183', 'Marbled_078', 'Matted_069', 'Mesh_114','Perforated_037','Stratified_154','Woven_001','Woven_068','Woven_104','Woven_125','Woven_127',
+    ]
+
+    def __init__(self, root=DTD_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = DTDSolver(root=DTD_ROOT)
+    runner.run()

data_preprocess/endo.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+ENDO_ROOT = os.path.join(DATA_ROOT, 'EndoTect')
+
+class ENDOSolver(object):
+    CLSNAMES = [
+        'endo',
+    ]
+
+    def __init__(self, root=ENDO_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = ENDOSolver(root=ENDO_ROOT)
+    runner.run()

data_preprocess/headct-pre.py

@@ -0,0 +1,41 @@
+import os
+import numpy as np
+from sklearn.model_selection import train_test_split
+import shutil
+import argparse
+
+from config import DATA_ROOT
+
+dataset_root = os.path.join(DATA_ROOT, 'head_ct')
+
+label_file = os.path.join(dataset_root, 'labels.csv')
+
+data = np.loadtxt(label_file, dtype=int, delimiter=',', skiprows=1)
+
+fnames = data[:, 0]
+label = data[:, 1]
+
+normal_fnames = fnames[label==0]
+outlier_fnames = fnames[label==1]
+
+
+target_root = '../datasets/HeadCT_anomaly_detection/headct'
+train_root = os.path.join(target_root, 'train/good')
+if not os.path.exists(train_root):
+    os.makedirs(train_root)
+
+test_normal_root = os.path.join(target_root, 'test/good')
+if not os.path.exists(test_normal_root):
+    os.makedirs(test_normal_root)
+for f in normal_fnames:
+    source = os.path.join(dataset_root, 'head_ct/', '{:0>3d}.png'.format(f))
+    shutil.copy(source, test_normal_root)
+
+test_outlier_root = os.path.join(target_root, 'test/defect')
+if not os.path.exists(test_outlier_root):
+    os.makedirs(test_outlier_root)
+for f in outlier_fnames:
+    source = os.path.join(dataset_root, 'head_ct/', '{:0>3d}.png'.format(f))
+    shutil.copy(source, test_outlier_root)
+
+print('Done')

data_preprocess/headct.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+# from dataset import MPDD_ROOT
+# from  dataset.mpdd import MPDD_ROOT
+
+
+HEADCT_ROOT = '../datasets/HeadCT_anomaly_detection'
+class HEADCTSolver(object):
+    CLSNAMES = [
+        'headct',
+    ]
+
+    def __init__(self, root=HEADCT_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    img_names.sort()
+
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = HEADCTSolver(root=HEADCT_ROOT)
+    runner.run()

data_preprocess/isic.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+ISIC_ROOT = os.path.join(DATA_ROOT, 'ISIC')
+
+class ISICSolver(object):
+    CLSNAMES = [
+        'isic',
+    ]
+
+    def __init__(self, root=ISIC_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = ISICSolver(root=ISIC_ROOT)
+    runner.run()

data_preprocess/mpdd.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+MPDD_ROOT = os.path.join(DATA_ROOT, 'MPDD')
+
+class MPDDSolver(object):
+    CLSNAMES = [
+        'bracket_black', 'bracket_brown', 'bracket_white', 'connector', 'metal_plate','tubes',
+    ]
+
+    def __init__(self, root=MPDD_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = MPDDSolver(root=MPDD_ROOT)
+    runner.run()

data_preprocess/mvtec.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from dataset import MVTEC_ROOT
+
+class MVTecSolver(object):
+    CLSNAMES = [
+        'bottle', 'cable', 'capsule', 'carpet', 'grid',
+        'hazelnut', 'leather', 'metal_nut', 'pill', 'screw',
+        'tile', 'toothbrush', 'transistor', 'wood', 'zipper',
+    ]
+
+    def __init__(self, root=MVTEC_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = MVTecSolver(root=MVTEC_ROOT)
+    runner.run()

data_preprocess/sdd-pre.py

@@ -0,0 +1,75 @@
+import os
+import numpy as np
+from sklearn.model_selection import train_test_split
+import cv2
+import argparse
+
+from config import DATA_ROOT
+
+dataset_root = os.path.join(DATA_ROOT, 'KolektorSDD')
+
+dirs = os.listdir(dataset_root)
+normal_images = list()
+normal_labels = list()
+normal_fname = list()
+outlier_images = list()
+outlier_labels = list()
+outlier_fname = list()
+for d in dirs:
+    files = os.listdir(os.path.join(dataset_root, d))
+    images = list()
+    for f in files:
+        if 'jpg' in f[-3:]:
+            images.append(f)
+
+    for image in images:
+        split_images = list()
+        split_labels = list()
+        image_name = image.split('.')[0]
+        image_data = cv2.imread(os.path.join(dataset_root, d, image))
+        label_data = cv2.imread(os.path.join(dataset_root, d, image_name + '_label.bmp'))
+        if image_data.shape != label_data.shape:
+            raise ValueError
+        image_length = image_data.shape[0]
+        split_images.append(image_data[:image_length // 3, :, :])
+        split_images.append(image_data[image_length // 3:image_length * 2 // 3, :, :])
+        split_images.append(image_data[image_length * 2 // 3:, :, :])
+        split_labels.append(label_data[:image_length // 3, :, :])
+        split_labels.append(label_data[image_length // 3:image_length * 2 // 3, :, :])
+        split_labels.append(label_data[image_length * 2 // 3:, :, :])
+        for i, (im, la) in enumerate(zip(split_images, split_labels)):
+            if np.max(la) != 0:
+                outlier_images.append(im)
+                outlier_labels.append(la)
+                outlier_fname.append(d + '_' + image_name + '_' + str(i))
+            else:
+                normal_images.append(im)
+                normal_labels.append(la)
+                normal_fname.append(d + '_' + image_name + '_' + str(i))
+
+normal_train, normal_test, normal_name_train, normal_name_test = train_test_split(normal_images, normal_fname, test_size=0.25, random_state=42)
+
+target_root = '../datasets/SDD_anomaly_detection/SDD'
+train_root = os.path.join(target_root, 'train/good')
+if not os.path.exists(train_root):
+    os.makedirs(train_root)
+for image, name in zip(normal_train, normal_name_train):
+    cv2.imwrite(os.path.join(train_root, name + '.png'), image)
+
+test_root = os.path.join(target_root, 'test/good')
+if not os.path.exists(test_root):
+    os.makedirs(test_root)
+for image, name in zip(normal_test, normal_name_test):
+    cv2.imwrite(os.path.join(test_root, name + '.png'), image)
+
+defect_root = os.path.join(target_root, 'test/defect')
+label_root = os.path.join(target_root, 'ground_truth/defect')
+if not os.path.exists(defect_root):
+    os.makedirs(defect_root)
+if not os.path.exists(label_root):
+    os.makedirs(label_root)
+for image, label, name in zip(outlier_images, outlier_labels, outlier_fname):
+    cv2.imwrite(os.path.join(defect_root, name + '.png'), image)
+    cv2.imwrite(os.path.join(label_root, name + '_mask.png'), label)
+
+print("Done")

data_preprocess/sdd.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+SDD_ROOT = os.path.join(DATA_ROOT, 'SDD_anomaly_detection')
+
+class SDDSolver(object):
+    CLSNAMES = [
+        'SDD',
+    ]
+
+    def __init__(self, root=SDD_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = SDDSolver(root=SDD_ROOT)
+    runner.run()

data_preprocess/tn3k.py

@@ -0,0 +1,52 @@
+import os
+import json
+import random
+from config import DATA_ROOT
+
+TN3K_ROOT = os.path.join(DATA_ROOT, 'TN3K')
+
+class TN3KSolver(object):
+    CLSNAMES = [
+        'tn3k',
+    ]
+
+    def __init__(self, root=TN3K_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        info = dict(train={}, test={})
+        for cls_name in self.CLSNAMES:
+            cls_dir = f'{self.root}/{cls_name}'
+            for phase in ['train', 'test']:
+                cls_info = []
+                species = os.listdir(f'{cls_dir}/{phase}')
+                for specie in species:
+                    is_abnormal = True if specie not in ['good'] else False
+                    img_names = os.listdir(f'{cls_dir}/{phase}/{specie}')
+                    mask_names = os.listdir(f'{cls_dir}/ground_truth/{specie}') if is_abnormal else None
+                    img_names.sort()
+                    mask_names.sort() if mask_names is not None else None
+                    for idx, img_name in enumerate(img_names):
+                        info_img = dict(
+                            img_path=f'{cls_name}/{phase}/{specie}/{img_name}',
+                            mask_path=f'{cls_name}/ground_truth/{specie}/{mask_names[idx]}' if is_abnormal else '',
+                            cls_name=cls_name,
+                            specie_name=specie,
+                            anomaly=1 if is_abnormal else 0,
+                        )
+                        cls_info.append(info_img)
+
+                info[phase][cls_name] = cls_info
+
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+if __name__ == '__main__':
+    runner = TN3KSolver(root=TN3K_ROOT)
+    runner.run()

data_preprocess/visa.py

@@ -0,0 +1,52 @@
+import os
+import json
+import pandas as pd
+import random
+from dataset import VISA_ROOT
+
+class VisASolver(object):
+    CLSNAMES = [
+        'candle', 'capsules', 'cashew', 'chewinggum', 'fryum',
+        'macaroni1', 'macaroni2', 'pcb1', 'pcb2', 'pcb3',
+        'pcb4', 'pipe_fryum',
+    ]
+
+    def __init__(self, root=VISA_ROOT, train_ratio=0.5):
+        self.root = root
+        self.meta_path = f'{root}/meta.json'
+        self.phases = ['train', 'test']
+        self.csv_data = pd.read_csv(f'{root}/split_csv/1cls.csv', header=0)
+        self.train_ratio = train_ratio
+
+    def run(self):
+        self.generate_meta_info()
+
+    def generate_meta_info(self):
+        columns = self.csv_data.columns  # [object, split, label, image, mask]
+        info = {phase: {} for phase in self.phases}
+        for cls_name in self.CLSNAMES:
+            cls_data = self.csv_data[self.csv_data[columns[0]] == cls_name]
+            for phase in self.phases:
+                cls_info = []
+                cls_data_phase = cls_data[cls_data[columns[1]] == phase]
+                cls_data_phase.index = list(range(len(cls_data_phase)))
+                for idx in range(cls_data_phase.shape[0]):
+                    data = cls_data_phase.loc[idx]
+                    is_abnormal = True if data[2] == 'anomaly' else False
+                    info_img = dict(
+                        img_path=data[3],
+                        mask_path=data[4] if is_abnormal else '',
+                        cls_name=cls_name,
+                        specie_name='',
+                        anomaly=1 if is_abnormal else 0,
+                    )
+                    cls_info.append(info_img)
+                info[phase][cls_name] = cls_info
+        with open(self.meta_path, 'w') as f:
+            f.write(json.dumps(info, indent=4) + "\n")
+
+
+
+if __name__ == '__main__':
+    runner = VisASolver(root=VISA_ROOT)
+    runner.run()

dataset/__init__.py

@@ -0,0 +1,68 @@
+from .mvtec import MVTEC_CLS_NAMES, MVTecDataset, MVTEC_ROOT
+from .visa import VISA_CLS_NAMES, VisaDataset, VISA_ROOT
+from .mpdd import MPDD_CLS_NAMES, MPDDDataset, MPDD_ROOT
+from .btad import BTAD_CLS_NAMES, BTADDataset, BTAD_ROOT
+from .sdd import SDD_CLS_NAMES, SDDDataset, SDD_ROOT
+from .dagm import DAGM_CLS_NAMES, DAGMDataset, DAGM_ROOT
+from .dtd import DTD_CLS_NAMES,DTDDataset,DTD_ROOT
+from .isic import ISIC_CLS_NAMES,ISICDataset,ISIC_ROOT
+from .colondb import ColonDB_CLS_NAMES, ColonDBDataset, ColonDB_ROOT
+from .clinicdb import ClinicDB_CLS_NAMES, ClinicDBDataset, ClinicDB_ROOT
+from .tn3k import TN3K_CLS_NAMES, TN3KDataset, TN3K_ROOT
+from .headct import HEADCT_CLS_NAMES,HEADCTDataset,HEADCT_ROOT
+from .brain_mri import BrainMRI_CLS_NAMES,BrainMRIDataset,BrainMRI_ROOT
+from .br35h import Br35h_CLS_NAMES,Br35hDataset,Br35h_ROOT
+from torch.utils.data import ConcatDataset
+
+dataset_dict = {
+    'br35h': (Br35h_CLS_NAMES, Br35hDataset, Br35h_ROOT),
+    'brain_mri': (BrainMRI_CLS_NAMES, BrainMRIDataset, BrainMRI_ROOT),
+    'btad': (BTAD_CLS_NAMES, BTADDataset, BTAD_ROOT),
+    'clinicdb': (ClinicDB_CLS_NAMES, ClinicDBDataset, ClinicDB_ROOT),
+    'colondb': (ColonDB_CLS_NAMES, ColonDBDataset, ColonDB_ROOT),
+    'dagm': (DAGM_CLS_NAMES, DAGMDataset, DAGM_ROOT),
+    'dtd': (DTD_CLS_NAMES, DTDDataset, DTD_ROOT),
+    'headct': (HEADCT_CLS_NAMES, HEADCTDataset, HEADCT_ROOT),
+    'isic': (ISIC_CLS_NAMES, ISICDataset, ISIC_ROOT),
+    'mpdd': (MPDD_CLS_NAMES, MPDDDataset, MPDD_ROOT),
+    'mvtec': (MVTEC_CLS_NAMES, MVTecDataset, MVTEC_ROOT),
+    'sdd': (SDD_CLS_NAMES, SDDDataset, SDD_ROOT),
+    'tn3k': (TN3K_CLS_NAMES, TN3KDataset, TN3K_ROOT),
+    'visa': (VISA_CLS_NAMES, VisaDataset, VISA_ROOT),
+}
+
+def get_data(dataset_type_list, transform, target_transform, training):
+    if not isinstance(dataset_type_list, list):
+        dataset_type_list = [dataset_type_list]
+
+    dataset_cls_names_list = []
+    dataset_instance_list = []
+    dataset_root_list = []
+    for dataset_type in dataset_type_list:
+        if dataset_dict.get(dataset_type, ''):
+            dataset_cls_names, dataset_instance, dataset_root = dataset_dict[dataset_type]
+            dataset_instance = dataset_instance(
+                clsnames=dataset_cls_names,
+                transform=transform,
+                target_transform=target_transform,
+                training=training
+            )
+
+            dataset_cls_names_list.append(dataset_cls_names)
+            dataset_instance_list.append(dataset_instance)
+            dataset_root_list.append(dataset_root)
+
+        else:
+            print(f'Only support {list(dataset_dict.keys())}, but entered {dataset_type}...')
+            raise NotImplementedError
+
+    if len(dataset_type_list) > 1:
+        dataset_instance = ConcatDataset(dataset_instance_list)
+        dataset_cls_names = dataset_cls_names_list
+        dataset_root = dataset_root_list
+    else:
+        dataset_instance = dataset_instance_list[0]
+        dataset_cls_names = dataset_cls_names_list[0]
+        dataset_root = dataset_root_list[0]
+
+    return dataset_cls_names, dataset_instance, dataset_root

dataset/base_dataset.py

@@ -0,0 +1,138 @@
+"""
+Base class for our zero-shot anomaly detection dataset
+"""
+import json
+import os
+import random
+import numpy as np
+import torch.utils.data as data
+from PIL import Image
+import cv2
+from config import DATA_ROOT
+
+
+class DataSolver:
+    def __init__(self, root, clsnames):
+        self.root = root
+        self.clsnames = clsnames
+        self.path = os.path.join(root, 'meta.json')
+
+    def run(self):
+        with open(self.path, 'r') as f:
+            info = json.load(f)
+
+        info_required = dict(train={}, test={})
+        for cls in self.clsnames:
+            for k in info.keys():
+                info_required[k][cls] = info[k][cls]
+
+        return info_required
+
+
+class BaseDataset(data.Dataset):
+    def __init__(self, clsnames, transform, target_transform, root, aug_rate=0., training=True):
+        self.root = root
+        self.transform = transform
+        self.target_transform = target_transform
+        self.aug_rate = aug_rate
+        self.training = training
+        self.data_all = []
+        self.cls_names = clsnames
+
+        solver = DataSolver(root, clsnames)
+        meta_info = solver.run()
+
+        self.meta_info = meta_info['test']  # Only utilize the test dataset for both training and testing
+        for cls_name in self.cls_names:
+            self.data_all.extend(self.meta_info[cls_name])
+
+        self.length = len(self.data_all)
+
+    def __len__(self):
+        return self.length
+
+    def combine_img(self, cls_name):
+        """
+        From April-GAN: https://github.com/ByChelsea/VAND-APRIL-GAN
+        Here we combine four images into a single image for data augmentation.
+        """
+        img_info = random.sample(self.meta_info[cls_name], 4)
+
+        img_ls = []
+        mask_ls = []
+
+        for data in img_info:
+            img_path = os.path.join(self.root, data['img_path'])
+            mask_path = os.path.join(self.root, data['mask_path'])
+
+            img = Image.open(img_path).convert('RGB')
+            img_ls.append(img)
+
+            if not data['anomaly']:
+                img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+            else:
+                img_mask = np.array(Image.open(mask_path).convert('L')) > 0
+                img_mask = Image.fromarray(img_mask.astype(np.uint8) * 255, mode='L')
+
+            mask_ls.append(img_mask)
+
+        # Image
+        image_width, image_height = img_ls[0].size
+        result_image = Image.new("RGB", (2 * image_width, 2 * image_height))
+        for i, img in enumerate(img_ls):
+            row = i // 2
+            col = i % 2
+            x = col * image_width
+            y = row * image_height
+            result_image.paste(img, (x, y))
+
+        # Mask
+        result_mask = Image.new("L", (2 * image_width, 2 * image_height))
+        for i, img in enumerate(mask_ls):
+            row = i // 2
+            col = i % 2
+            x = col * image_width
+            y = row * image_height
+            result_mask.paste(img, (x, y))
+
+        return result_image, result_mask
+
+    def __getitem__(self, index):
+        data = self.data_all[index]
+        img_path = os.path.join(self.root, data['img_path'])
+        mask_path = os.path.join(self.root, data['mask_path'])
+        cls_name = data['cls_name']
+        anomaly = data['anomaly']
+        random_number = random.random()
+
+        if self.training and random_number < self.aug_rate:
+            img, img_mask = self.combine_img(cls_name)
+        else:
+            if img_path.endswith('.tif'):
+                img = cv2.imread(img_path)
+                img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+            else:
+                img = Image.open(img_path).convert('RGB')
+            if anomaly == 0:
+                img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+            else:
+                if data['mask_path']:
+                    img_mask = np.array(Image.open(mask_path).convert('L')) > 0
+                    img_mask = Image.fromarray(img_mask.astype(np.uint8) * 255, mode='L')
+                else:
+                    img_mask = Image.fromarray(np.zeros((img.size[0], img.size[1])), mode='L')
+        # Transforms
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.target_transform is not None and img_mask is not None:
+            img_mask = self.target_transform(img_mask)
+        if img_mask is None:
+            img_mask = []
+
+        return {
+            'img': img,
+            'img_mask': img_mask,
+            'cls_name': cls_name,
+            'anomaly': anomaly,
+            'img_path': img_path
+        }

dataset/br35h.py

@@ -0,0 +1,18 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://www.kaggle.com/datasets/ahmedhamada0/brain-tumor-detection'''
+
+Br35h_CLS_NAMES = [
+    'br35h',
+]
+Br35h_ROOT = os.path.join(DATA_ROOT, 'Br35h_anomaly_detection')
+
+class Br35hDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=Br35h_CLS_NAMES, aug_rate=0.0, root=Br35h_ROOT, training=True):
+        super(Br35hDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+