initial commit

Dockerfile

@@ -0,0 +1,14 @@
+# -----------------------------------------------------------------------------
+#  A sample Dockerfile to help you replicate our test environment
+# -----------------------------------------------------------------------------
+
+FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime
+WORKDIR /app
+COPY . .
+
+# Install your python and apt requirements
+RUN pip install -r requirements.txt
+RUN apt-get update && apt-get install $(cat apt_requirements.txt) -y
+RUN chmod +x run.sh
+
+CMD ["python3", "runner.py"]

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

@@ -0,0 +1,189 @@
+# AdaCLIP (Detecting Anomalies for Novel Categories)
+[![HuggingFace Space](https://img.shields.io/badge/🤗-HuggingFace%20Space-cyan.svg)](https://huggingface.co/spaces/Caoyunkang/AdaCLIP)
+
+> [**ECCV 24**] [**AdaCLIP: Adapting CLIP with Hybrid Learnable Prompts for Zero-Shot Anomaly Detection**](https://arxiv.org/abs/2407.15795).
+>
+> 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.
+
+## Corrections
+- The description to the utilized training set in our paper is not accurate. By default, we utilize MVTec AD & ColonDB for training,
+and VisA & ClinicDB are utilized for evaluations on MVTec AD & ColonDB. 
+
+## 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](https://drive.google.com/file/d/12IukAqxOj497J4F0Mel-FvaONM030qwP/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1k36IMP4w32hY9BXOUM5ZmA?pwd=kxud) | Anomaly Detection & Localization |
+| VisA    | [Google Drive](https://drive.google.com/file/d/1U0MZVro5yGgaHNQ8kWb3U1a0Qlz4HiHI/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/15CIsP-ulZ1AN0_3quA068w?pwd=lmgc) | Anomaly Detection & Localization |
+| MPDD    | [Google Drive](https://drive.google.com/file/d/1cLkZs8pN8onQzfyNskeU_836JLjrtJz1/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/11T3mkloDCl7Hze5znkXOQA?pwd=4p7m) | Anomaly Detection & Localization | 
+| BTAD    | [Google Drive](https://drive.google.com/file/d/19Kd8jJLxZExwiTc9__6_r_jPqkmTXt4h/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1f4Tq-EXRz6iAswygH2WbFg?pwd=a60n) | Anomaly Detection & Localization |
+| KSDD    | [Google Drive](https://drive.google.com/file/d/13UidsM1taqEAVV_JJTBiCV1D3KUBpmpj/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/12EaOdkSbdK85WX5ajrfjQw?pwd=6n3z) | Anomaly Detection & Localization |
+| DAGM    | [Google Drive](https://drive.google.com/file/d/1f4sm8hpWQRzZMpvM-j7Q3xPG2vtdwvTy/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1JpDUJIksD99t003dNF1y9g?pwd=u3aq) | Anomaly Detection & Localization |
+| DTD-Synthetic    | [Google Drive](https://drive.google.com/file/d/1em51XXz5_aBNRJlJxxv3-Ed1dO9H3QgS/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/16FlvIBWtjaDzWxlZfWjNeg?pwd=aq5c) | Anomaly Detection & Localization |
+
+
+
+
+#### Medical Visual Anomaly Detection Datasets
+| Dataset | Google Drive | Baidu Drive | Task
+|------------|------------------|------------------|  ------------------|
+| HeadCT    | [Google Drive](https://drive.google.com/file/d/1ore0yCV31oLwwC--YUuTQfij-f2V32O2/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/16PfXWJlh6Y9vkecY9IownA?pwd=svsl) | Anomaly Detection |
+| BrainMRI    | [Google Drive](https://drive.google.com/file/d/1JLYyzcPG3ULY2J_aw1SY9esNujYm9GKd/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1UgGlTR-ABWAEiVUX-QSPhA?pwd=vh9e) | Anomaly Detection |
+| Br35H    | [Google Drive](https://drive.google.com/file/d/1qaZ6VJDRk3Ix3oVp3NpFyTsqXLJ_JjQy/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1yCS6t3ht6qwJgM06YsU3mg?pwd=ps1e) | Anomaly Detection |
+| ISIC    | [Google Drive](https://drive.google.com/file/d/1atZwmnFsz7mCsHWBZ8pkL_-Eul9bKFEx/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1Mf0w8RFY9ECZBEoNTyV3ZA?pwd=p954) | Anomaly Localization |
+| ColonDB    | [Google Drive](https://drive.google.com/file/d/1tjZ0o5dgzka3wf_p4ErSRJ9fcC-RJK8R/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1nJ4L65vfNFGpkK_OJjLoVg?pwd=v8q7) | Anomaly Localization |
+| ClinicDB    | [Google Drive](https://drive.google.com/file/d/1ciqZwMs1smSGDlwQ6tsr6YzylrqQBn9n/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1TPysfqhA_sXRPLGNwWBX6Q?pwd=3da6) | Anomaly Localization |
+| TN3K    | [Google Drive](https://drive.google.com/file/d/1LuKEMhrUGwFBlGCaej46WoooH89V3O8_/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1i5jMofCcRFcUdteq8VMEOQ?pwd=aoez) | 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](https://pan.baidu.com/s/1K9JhNAmmDt4n5Sqlq4-5hQ?pwd=fks1) | 
+| VisA & ColonDB    | [Google Drive](https://drive.google.com/file/d/1QGmPB0ByPZQ7FucvGODMSz7r5Ke5wx9W/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1GmRCylpboPseT9lguCO9nw?pwd=fvvf) | 
+| All Datasets Mentioned Above   | [Google Drive](https://drive.google.com/file/d/1Cgkfx3GAaSYnXPLolx-P7pFqYV0IVzZF/view?usp=drive_link) | [Baidu Drive](https://pan.baidu.com/s/1J4aFAOhUbeYOBfZFbkOixA?pwd=0ts3) |
+
+
+### Train
+
+By default, we use MVTec AD & Colondb 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 & Colondb, we use VisA & ClinicDB 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
+```
+
+Or visit our [Online Demo](https://huggingface.co/spaces/Caoyunkang/AdaCLIP) for a quick start. The three pre-trained weights mentioned are available there. Feel free to test them with your own data!
+
+Please note that we currently do not have a GPU environment for our Hugging Face Space, so inference for a single image may take approximately 50 seconds.
+
+![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
+
+@inproceedings{AdaCLIP,
+  title={AdaCLIP: Adapting CLIP with Hybrid Learnable Prompts for Zero-Shot Anomaly Detection},
+  author={Cao, Yunkang and Zhang, Jiangning and Frittoli, Luca and Cheng, Yuqi and Shen, Weiming and Boracchi, Giacomo},
+  booktitle={European Conference on Computer Vision},
+  year={2024}
+}
+
+```

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 = 'data'

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
+        )
+

dataset/brain_mri.py

@@ -0,0 +1,16 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://www.kaggle.com/datasets/navoneel/brain-mri-images-for-brain-tumor-detection'''
+BrainMRI_CLS_NAMES = [
+    'brain_mri',
+]
+BrainMRI_ROOT = os.path.join(DATA_ROOT, 'BrainMRI')
+
+class BrainMRIDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=BrainMRI_CLS_NAMES, aug_rate=0.0, root=BrainMRI_ROOT, training=True):
+        super(BrainMRIDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )

dataset/btad.py

@@ -0,0 +1,16 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://avires.dimi.uniud.it/papers/btad/btad.zip'''
+BTAD_CLS_NAMES = [
+    '01', '02', '03',
+]
+BTAD_ROOT = os.path.join(DATA_ROOT, 'BTech_Dataset_transformed')
+
+class BTADDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=BTAD_CLS_NAMES, aug_rate=0.0, root=BTAD_ROOT, training=True):
+        super(BTADDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )

dataset/clinicdb.py

@@ -0,0 +1,16 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://paperswithcode.com/dataset/cvc-clinicdb'''
+ClinicDB_CLS_NAMES = [
+    'ClinicDB',
+]
+ClinicDB_ROOT = os.path.join(DATA_ROOT, 'CVC-ClinicDB')
+
+class ClinicDBDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=ClinicDB_CLS_NAMES, aug_rate=0.0, root=ClinicDB_ROOT, training=True):
+        super(ClinicDBDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )

dataset/colondb.py

@@ -0,0 +1,18 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: http://mv.cvc.uab.es/projects/colon-qa/cvccolondb'''
+ColonDB_CLS_NAMES = [
+    'ColonDB',
+]
+ColonDB_ROOT = os.path.join(DATA_ROOT, 'CVC-ColonDB')
+
+class ColonDBDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=ColonDB_CLS_NAMES, aug_rate=0.0, root=ColonDB_ROOT, training=True):
+        super(ColonDBDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+
+

dataset/dagm.py

@@ -0,0 +1,16 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://hci.iwr.uni-heidelberg.de/content/weakly-supervised-learning-industrial-optical-inspection'''
+DAGM_CLS_NAMES = [
+    'Class1', 'Class2', 'Class3', 'Class4', 'Class5','Class6','Class7','Class8','Class9','Class10',
+]
+DAGM_ROOT = os.path.join(DATA_ROOT, 'DAGM_anomaly_detection')
+
+class DAGMDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=DAGM_CLS_NAMES, aug_rate=0.0, root=DAGM_ROOT, training=True):
+        super(DAGMDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )

dataset/dtd.py

@@ -0,0 +1,16 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://drive.google.com/drive/folders/10OyPzvI3H6llCZBxKxFlKWt1Pw1tkMK1'''
+DTD_CLS_NAMES = [
+    'Blotchy_099', 'Fibrous_183', 'Marbled_078', 'Matted_069', 'Mesh_114','Perforated_037','Stratified_154','Woven_001','Woven_068','Woven_104','Woven_125','Woven_127',
+]
+DTD_ROOT = os.path.join(DATA_ROOT, 'DTD-Synthetic')
+
+class DTDDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=DTD_CLS_NAMES, aug_rate=0.0, root=DTD_ROOT, training=True):
+        super(DTDDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )

dataset/headct.py

@@ -0,0 +1,18 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://www.kaggle.com/datasets/felipekitamura/head-ct-hemorrhage'''
+HEADCT_CLS_NAMES = [
+    'headct',
+]
+HEADCT_ROOT = os.path.join(DATA_ROOT, 'HeadCT_anomaly_detection')
+
+class HEADCTDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=HEADCT_CLS_NAMES, aug_rate=0.0, root=HEADCT_ROOT, training=True):
+        super(HEADCTDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+
+

dataset/isic.py

@@ -0,0 +1,18 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://challenge.isic-archive.com/data/'''
+ISIC_CLS_NAMES = [
+    'isic',
+]
+ISIC_ROOT = os.path.join(DATA_ROOT, 'ISIC')
+
+class ISICDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=ISIC_CLS_NAMES, aug_rate=0.0, root=ISIC_ROOT, training=True):
+        super(ISICDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+
+

dataset/mpdd.py

@@ -0,0 +1,17 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://github.com/stepanje/MPDD'''
+MPDD_CLS_NAMES = [
+    'bracket_black', 'bracket_brown', 'bracket_white', 'connector', 'metal_plate','tubes',
+]
+MPDD_ROOT = os.path.join(DATA_ROOT, 'MPDD')
+
+class MPDDDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=MPDD_CLS_NAMES, aug_rate=0.0, root=MPDD_ROOT, training=True):
+        super(MPDDDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+

dataset/mvtec.py

@@ -0,0 +1,19 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://paperswithcode.com/dataset/mvtecad'''
+
+MVTEC_CLS_NAMES = [
+    'bottle', 'cable', 'capsule', 'carpet', 'grid',
+    'hazelnut', 'leather', 'metal_nut', 'pill', 'screw',
+    'tile', 'toothbrush', 'transistor', 'wood', 'zipper',
+]
+MVTEC_ROOT = os.path.join(DATA_ROOT, 'mvtec_anomaly_detection')
+
+class MVTecDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=MVTEC_CLS_NAMES, aug_rate=0.2, root=MVTEC_ROOT, training=True):
+        super(MVTecDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )

dataset/sdd.py

@@ -0,0 +1,18 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://data.vicos.si/datasets/KSDD/KolektorSDD.zip'''
+SDD_CLS_NAMES = [
+    'SDD',
+]
+SDD_ROOT = os.path.join(DATA_ROOT, 'SDD_anomaly_detection')
+
+
+class SDDDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=SDD_CLS_NAMES, aug_rate=0.0, root=SDD_ROOT, training=True):
+        super(SDDDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+

dataset/tn3k.py

@@ -0,0 +1,18 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://ieeexplore.ieee.org/document/9434087/references#references'''
+TN3K_CLS_NAMES = [
+    'tn3k',
+]
+TN3K_ROOT = os.path.join(DATA_ROOT, 'TN3K')
+
+class TN3KDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=TN3K_CLS_NAMES, aug_rate=0.0, root=TN3K_ROOT, training=True):
+        super(TN3KDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+
+

dataset/visa.py

@@ -0,0 +1,20 @@
+import os
+from .base_dataset import BaseDataset
+from config import DATA_ROOT
+
+'''dataset source: https://amazon-visual-anomaly.s3.us-west-2.amazonaws.com/VisA_20220922.tar'''
+VISA_CLS_NAMES = [
+    'candle', 'capsules', 'cashew', 'chewinggum', 'fryum',
+    'macaroni1', 'macaroni2', 'pcb1', 'pcb2', 'pcb3',
+    'pcb4', 'pipe_fryum',
+]
+
+VISA_ROOT = os.path.join(DATA_ROOT, 'VisA_20220922')
+
+class VisaDataset(BaseDataset):
+    def __init__(self, transform, target_transform, clsnames=VISA_CLS_NAMES, aug_rate=0.0, root=VISA_ROOT, training=True):
+        super(VisaDataset, self).__init__(
+            clsnames=clsnames, transform=transform, target_transform=target_transform,
+            root=root, aug_rate=aug_rate, training=training
+        )
+

datasets/rayan_dataset.py

@@ -0,0 +1,127 @@
+# -----------------------------------------------------------------------------
+#  Do Not Alter This File!
+# -----------------------------------------------------------------------------
+#  The following code is part of the logic used for loading and evaluating your
+#  output scores. Please DO NOT modify this section, as upon your submission,
+#  the whole evaluation logic will be overwritten by the original code.
+# -----------------------------------------------------------------------------
+#  If you'd like to make modifications, you can create a completely new Dataset
+#  class or a child class that inherits from this one and use that with your
+#  data loader.
+# -----------------------------------------------------------------------------
+
+import os
+from enum import Enum
+
+import PIL
+import torch
+from torchvision import transforms
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
+class DatasetSplit(Enum):
+    TRAIN = "train"
+    VAL = "val"
+    TEST = "test"
+
+
+class RayanDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        source,
+        classname,
+        input_size=518,
+        output_size=224,
+        split=DatasetSplit.TEST,
+        external_transform=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.source = source
+        self.split = split
+        self.classnames_to_use = [classname]
+        self.imgpaths_per_class, self.data_to_iterate = self.get_image_data()
+
+        if external_transform is None:
+            self.transform_img = [
+                transforms.Resize((input_size, input_size)),
+                transforms.CenterCrop(input_size),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+            ]
+            self.transform_img = transforms.Compose(self.transform_img)
+        else:
+            self.transform_img = external_transform
+
+        # Output size of the mask has to be of shape: 1×224×224
+        self.transform_mask = [
+            transforms.Resize((output_size, output_size)),
+            transforms.CenterCrop(output_size),
+            transforms.ToTensor(),
+        ]
+        self.transform_mask = transforms.Compose(self.transform_mask)
+        self.output_shape = (1, output_size, output_size)
+
+    def __getitem__(self, idx):
+        classname, anomaly, image_path, mask_path = self.data_to_iterate[idx]
+        image = PIL.Image.open(image_path).convert("RGB")
+        image = self.transform_img(image)
+
+        if self.split == DatasetSplit.TEST and mask_path is not None:
+            mask = PIL.Image.open(mask_path).convert("L")
+            mask = self.transform_mask(mask) > 0
+        else:
+            mask = torch.zeros([*self.output_shape])
+
+        return {
+            "image": image,
+            "mask": mask,
+            "is_anomaly": int(anomaly != "good"),
+            "image_path": image_path,
+        }
+
+    def __len__(self):
+        return len(self.data_to_iterate)
+
+    def get_image_data(self):
+        imgpaths_per_class = {}
+        maskpaths_per_class = {}
+
+        for classname in self.classnames_to_use:
+            classpath = os.path.join(self.source, classname, self.split.value)
+            maskpath = os.path.join(self.source, classname, "ground_truth")
+            anomaly_types = os.listdir(classpath)
+
+            imgpaths_per_class[classname] = {}
+            maskpaths_per_class[classname] = {}
+
+            for anomaly in anomaly_types:
+                anomaly_path = os.path.join(classpath, anomaly)
+                anomaly_files = sorted(os.listdir(anomaly_path))
+                imgpaths_per_class[classname][anomaly] = [
+                    os.path.join(anomaly_path, x) for x in anomaly_files
+                ]
+
+                if self.split == DatasetSplit.TEST and anomaly != "good":
+                    anomaly_mask_path = os.path.join(maskpath, anomaly)
+                    anomaly_mask_files = sorted(os.listdir(anomaly_mask_path))
+                    maskpaths_per_class[classname][anomaly] = [
+                        os.path.join(anomaly_mask_path, x) for x in anomaly_mask_files
+                    ]
+                else:
+                    maskpaths_per_class[classname]["good"] = None
+
+        data_to_iterate = []
+        for classname in sorted(imgpaths_per_class.keys()):
+            for anomaly in sorted(imgpaths_per_class[classname].keys()):
+                for i, image_path in enumerate(imgpaths_per_class[classname][anomaly]):
+                    data_tuple = [classname, anomaly, image_path]
+                    if self.split == DatasetSplit.TEST and anomaly != "good":
+                        data_tuple.append(maskpaths_per_class[classname][anomaly][i])
+                    else:
+                        data_tuple.append(None)
+                    data_to_iterate.append(data_tuple)
+
+        return imgpaths_per_class, data_to_iterate

docker-compose.yml

@@ -0,0 +1,21 @@
+# -----------------------------------------------------------------------------
+#  A sample Docker Compose file to help you replicate our test environment
+# -----------------------------------------------------------------------------
+
+services:
+    zsad-service:
+        image: zsad-image:1
+        build:
+            context: .
+        container_name: zsad-container
+        volumes:
+            - ./shared_folder:/app/output
+        deploy:
+            resources:
+                reservations:
+                    devices:
+                        - driver: nvidia
+                          count: all
+                          capabilities: [gpu]
+
+        command: [ "python3", "runner.py" ]

evaluation/base_eval.py

@@ -0,0 +1,293 @@
+# -----------------------------------------------------------------------------
+#  Do Not Alter This File!
+# -----------------------------------------------------------------------------
+#  The following code is part of the logic used for loading and evaluating your
+#  output scores. Please DO NOT modify this section, as upon your submission,
+#  the whole evaluation logic will be overwritten by the original code.
+# -----------------------------------------------------------------------------
+
+import warnings
+import os
+from pathlib import Path
+import csv
+import json
+import torch
+
+import datasets.rayan_dataset as rayan_dataset
+from evaluation.utils.metrics import compute_metrics
+
+warnings.filterwarnings("ignore")
+
+
+class BaseEval:
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.device = torch.device(
+            "cuda:{}".format(cfg["device"]) if torch.cuda.is_available() else "cpu"
+        )
+
+        self.path = cfg["datasets"]["data_path"]
+        self.dataset = cfg["datasets"]["dataset_name"]
+        self.save_csv = cfg["testing"]["save_csv"]
+        self.save_json = cfg["testing"]["save_json"]
+        self.categories = cfg["datasets"]["class_name"]
+        if isinstance(self.categories, str):
+            if self.categories.lower() == "all":
+                if self.dataset == "rayan_dataset":
+                    self.categories = self.get_available_class_names(self.path)
+            else:
+                self.categories = [self.categories]
+        self.output_dir = cfg["testing"]["output_dir"]
+        os.makedirs(self.output_dir, exist_ok=True)
+        self.scores_dir = cfg["testing"]["output_scores_dir"]
+        self.class_name_mapping_dir = cfg["testing"]["class_name_mapping_dir"]
+
+        self.leaderboard_metric_weights = {
+            "image_auroc": 1.2,
+            "image_ap": 1.1,
+            "image_f1": 1.1,
+            "pixel_auroc": 1.0,
+            "pixel_aupro": 1.4,
+            "pixel_ap": 1.3,
+            "pixel_f1": 1.3,
+        }
+
+    def get_available_class_names(self, root_data_path):
+        all_items = os.listdir(root_data_path)
+        folder_names = [
+            item
+            for item in all_items
+            if os.path.isdir(os.path.join(root_data_path, item))
+        ]
+
+        return folder_names
+
+    def load_datasets(self, category):
+        dataset_classes = {
+            "rayan_dataset": rayan_dataset.RayanDataset,
+        }
+
+        dataset_splits = {
+            "rayan_dataset": rayan_dataset.DatasetSplit.TEST,
+        }
+
+        test_dataset = dataset_classes[self.dataset](
+            source=self.path,
+            split=dataset_splits[self.dataset],
+            classname=category,
+        )
+        return test_dataset
+
+    def get_category_metrics(self, category):
+        print(f"Loading scores of '{category}'")
+        gt_sp, pr_sp, gt_px, pr_px, _ = self.load_category_scores(category)
+
+        print(f"Computing metrics for '{category}'")
+        image_metric, pixel_metric = compute_metrics(gt_sp, pr_sp, gt_px, pr_px)
+
+        return image_metric, pixel_metric
+
+    def load_category_scores(self, category):
+        raise NotImplementedError()
+
+    def get_scores_path_for_image(self, image_path):
+        """example image_path: './data/photovoltaic_module/test/good/037.png'"""
+        path = Path(image_path)
+
+        category, split, anomaly_type = path.parts[-4:-1]
+        image_name = path.stem
+
+        return os.path.join(
+            self.scores_dir, category, split, anomaly_type, f"{image_name}_scores.json"
+        )
+
+    def calc_leaderboard_score(self, **metrics):
+        weighted_sum = 0
+        total_weight = 0
+        for key, weight in self.leaderboard_metric_weights.items():
+            metric = metrics.get(key)
+            weighted_sum += metric * weight
+            total_weight += weight
+
+        if total_weight == 0:
+            return 0
+
+        return weighted_sum / total_weight
+
+    def main(self):
+        image_auroc_list = []
+        image_f1_list = []
+        image_ap_list = []
+        pixel_auroc_list = []
+        pixel_f1_list = []
+        pixel_ap_list = []
+        pixel_aupro_list = []
+        leaderboard_score_list = []
+        for category in self.categories:
+            image_metric, pixel_metric = self.get_category_metrics(
+                category=category,
+            )
+            image_auroc, image_f1, image_ap = image_metric
+            pixel_auroc, pixel_f1, pixel_ap, pixel_aupro = pixel_metric
+            leaderboard_score = self.calc_leaderboard_score(
+                image_auroc=image_auroc,
+                image_f1=image_f1,
+                image_ap=image_ap,
+                pixel_auroc=pixel_auroc,
+                pixel_aupro=pixel_aupro,
+                pixel_f1=pixel_f1,
+                pixel_ap=pixel_ap,
+            )
+
+            image_auroc_list.append(image_auroc)
+            image_f1_list.append(image_f1)
+            image_ap_list.append(image_ap)
+            pixel_auroc_list.append(pixel_auroc)
+            pixel_f1_list.append(pixel_f1)
+            pixel_ap_list.append(pixel_ap)
+            pixel_aupro_list.append(pixel_aupro)
+            leaderboard_score_list.append(leaderboard_score)
+
+            print(category)
+            print(
+                "[image level] auroc:{}, f1:{}, ap:{}".format(
+                    image_auroc * 100,
+                    image_f1 * 100,
+                    image_ap * 100,
+                )
+            )
+            print(
+                "[pixel level] auroc:{}, f1:{}, ap:{}, aupro:{}".format(
+                    pixel_auroc * 100,
+                    pixel_f1 * 100,
+                    pixel_ap * 100,
+                    pixel_aupro * 100,
+                )
+            )
+            print(
+                "leaderboard score:{}".format(
+                    leaderboard_score * 100,
+                )
+            )
+
+        image_auroc_mean = sum(image_auroc_list) / len(image_auroc_list)
+        image_f1_mean = sum(image_f1_list) / len(image_f1_list)
+        image_ap_mean = sum(image_ap_list) / len(image_ap_list)
+        pixel_auroc_mean = sum(pixel_auroc_list) / len(pixel_auroc_list)
+        pixel_f1_mean = sum(pixel_f1_list) / len(pixel_f1_list)
+        pixel_ap_mean = sum(pixel_ap_list) / len(pixel_ap_list)
+        pixel_aupro_mean = sum(pixel_aupro_list) / len(pixel_aupro_list)
+        leaderboard_score_mean = sum(leaderboard_score_list) / len(
+            leaderboard_score_list
+        )
+
+        print("mean")
+        print(
+            "[image level] auroc:{}, f1:{}, ap:{}".format(
+                image_auroc_mean * 100, image_f1_mean * 100, image_ap_mean * 100
+            )
+        )
+        print(
+            "[pixel level] auroc:{}, f1:{}, ap:{}, aupro:{}".format(
+                pixel_auroc_mean * 100,
+                pixel_f1_mean * 100,
+                pixel_ap_mean * 100,
+                pixel_aupro_mean * 100,
+            )
+        )
+        print(
+            "leaderboard score:{}".format(
+                leaderboard_score_mean * 100,
+            )
+        )
+
+        # Save the final results as a csv file
+        if self.save_csv:
+            with open(self.class_name_mapping_dir, "r") as f:
+                class_name_mapping_dict = json.load(f)
+            csv_data = [
+                [
+                    "Category",
+                    "pixel_auroc",
+                    "pixel_f1",
+                    "pixel_ap",
+                    "pixel_aupro",
+                    "image_auroc",
+                    "image_f1",
+                    "image_ap",
+                    "leaderboard_score",
+                ]
+            ]
+            for i, category in enumerate(self.categories):
+                csv_data.append(
+                    [
+                        class_name_mapping_dict[category],
+                        pixel_auroc_list[i] * 100,
+                        pixel_f1_list[i] * 100,
+                        pixel_ap_list[i] * 100,
+                        pixel_aupro_list[i] * 100,
+                        image_auroc_list[i] * 100,
+                        image_f1_list[i] * 100,
+                        image_ap_list[i] * 100,
+                        leaderboard_score_list[i] * 100,
+                    ]
+                )
+            csv_data.append(
+                [
+                    "mean",
+                    pixel_auroc_mean * 100,
+                    pixel_f1_mean * 100,
+                    pixel_ap_mean * 100,
+                    pixel_aupro_mean * 100,
+                    image_auroc_mean * 100,
+                    image_f1_mean * 100,
+                    image_ap_mean * 100,
+                    leaderboard_score_mean * 100,
+                ]
+            )
+
+            csv_file_path = os.path.join(self.output_dir, "results.csv")
+            with open(csv_file_path, mode="w", newline="") as file:
+                writer = csv.writer(file)
+                writer.writerows(csv_data)
+
+        # Save the final results as a json file
+        if self.save_json:
+            json_data = []
+            with open(self.class_name_mapping_dir, "r") as f:
+                class_name_mapping_dict = json.load(f)
+            for i, category in enumerate(self.categories):
+                json_data.append(
+                    {
+                        "Category": class_name_mapping_dict[category],
+                        "pixel_auroc": pixel_auroc_list[i] * 100,
+                        "pixel_f1": pixel_f1_list[i] * 100,
+                        "pixel_ap": pixel_ap_list[i] * 100,
+                        "pixel_aupro": pixel_aupro_list[i] * 100,
+                        "image_auroc": image_auroc_list[i] * 100,
+                        "image_f1": image_f1_list[i] * 100,
+                        "image_ap": image_ap_list[i] * 100,
+                        "leaderboard_score": leaderboard_score_list[i] * 100,
+                    }
+                )
+            json_data.append(
+                {
+                    "Category": "mean",
+                    "pixel_auroc": pixel_auroc_mean * 100,
+                    "pixel_f1": pixel_f1_mean * 100,
+                    "pixel_ap": pixel_ap_mean * 100,
+                    "pixel_aupro": pixel_aupro_mean * 100,
+                    "image_auroc": image_auroc_mean * 100,
+                    "image_f1": image_f1_mean * 100,
+                    "image_ap": image_ap_mean * 100,
+                    "leaderboard_score": leaderboard_score_mean * 100,
+                }
+            )
+
+            json_file_path = os.path.join(self.output_dir, "results.json")
+            with open(json_file_path, mode="w") as file:
+                final_json = {
+                    "result": leaderboard_score_mean * 100,
+                    "metadata": json_data,
+                }
+                json.dump(final_json, file, indent=4)

evaluation/class_name_mapping.json

@@ -0,0 +1,5 @@
+{
+    "pill": "industrial_01",
+    "photovoltaic_module": "industrial_02",
+    "capsules": "industrial_03"
+}

evaluation/eval_main.py

@@ -0,0 +1,78 @@
+# -----------------------------------------------------------------------------
+#  Do Not Alter This File!
+# -----------------------------------------------------------------------------
+#  The following code is part of the logic used for loading and evaluating your
+#  output scores. Please DO NOT modify this section, as upon your submission,
+#  the whole evaluation logic will be overwritten by the original code.
+# -----------------------------------------------------------------------------
+
+import warnings
+import argparse
+import os
+import sys
+
+sys.path.append(os.getcwd())
+from evaluation.json_score import JsonScoreEvaluator
+
+warnings.filterwarnings("ignore")
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Rayan ZSAD Evaluation Code")
+    parser.add_argument("--data_path", type=str, default=None, help="dataset path")
+    parser.add_argument("--dataset_name", type=str, default=None, help="dataset name")
+    parser.add_argument("--class_name", type=str, default=None, help="category")
+    parser.add_argument("--device", type=int, default=None, help="gpu id")
+    parser.add_argument(
+        "--output_dir", type=str, default=None, help="save results path"
+    )
+    parser.add_argument(
+        "--output_scores_dir", type=str, default=None, help="save scores path"
+    )
+    parser.add_argument("--save_csv", type=str, default=None, help="save csv")
+    parser.add_argument("--save_json", type=str, default=None, help="save json")
+
+    parser.add_argument(
+        "--class_name_mapping_dir",
+        type=str,
+        default=None,
+        help="mapping from actual class names to class numbers",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def load_args(cfg, args):
+    cfg["datasets"]["data_path"] = args.data_path
+    assert os.path.exists(
+        cfg["datasets"]["data_path"]
+    ), f"The dataset path {cfg['datasets']['data_path']} does not exist."
+    cfg["datasets"]["dataset_name"] = args.dataset_name
+    cfg["datasets"]["class_name"] = args.class_name
+    cfg["device"] = args.device
+    if isinstance(cfg["device"], int):
+        cfg["device"] = str(cfg["device"])
+    cfg["testing"]["output_dir"] = args.output_dir
+    cfg["testing"]["output_scores_dir"] = args.output_scores_dir
+    os.makedirs(cfg["testing"]["output_scores_dir"], exist_ok=True)
+
+    cfg["testing"]["class_name_mapping_dir"] = args.class_name_mapping_dir
+    if args.save_csv.lower() == "true":
+        cfg["testing"]["save_csv"] = True
+    else:
+        cfg["testing"]["save_csv"] = False
+
+    if args.save_json.lower() == "true":
+        cfg["testing"]["save_json"] = True
+    else:
+        cfg["testing"]["save_json"] = False
+
+    return cfg
+
+
+if __name__ == "__main__":
+    args = get_args()
+    cfg = load_args(cfg={"datasets": {}, "testing": {}, "models": {}}, args=args)
+    print(cfg)
+    model = JsonScoreEvaluator(cfg=cfg)
+    model.main()

evaluation/json_score.py

@@ -0,0 +1,98 @@
+# -----------------------------------------------------------------------------
+#  Do Not Alter This File!
+# -----------------------------------------------------------------------------
+#  The following code is part of the logic used for loading and evaluating your
+#  output scores. Please DO NOT modify this section, as upon your submission,
+#  the whole evaluation logic will be overwritten by the original code.
+# -----------------------------------------------------------------------------
+
+import warnings
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from evaluation.base_eval import BaseEval
+from evaluation.utils.json_helpers import json_to_dict
+
+warnings.filterwarnings("ignore")
+
+
+class JsonScoreEvaluator(BaseEval):
+    """
+    Evaluates anomaly detection performance based on pre-computed scores stored in JSON files.
+
+    This class extends the BaseEval class and specializes in reading scores from JSON files,
+    computing evaluation metrics, and optionally saving results to CSV or JSON format.
+
+    Notes:
+        - Score files are expected to follow the exact dataset structure.
+                `{category}/{split}/{anomaly_type}/{image_name}_scores.json`
+          e.g., `photovoltaic_module/test/good/037_scores.json`
+        - Score files are expected to be at `self.scores_dir`.
+
+    Example usage:
+        >>> evaluator = JsonScoreEvaluator(cfg)
+        >>> results = evaluator.main()
+    """
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def get_scores_for_image(self, image_path):
+        image_scores_path = self.get_scores_path_for_image(image_path)
+        image_scores = json_to_dict(image_scores_path)
+
+        return image_scores
+
+    def load_category_scores(self, category):
+        cls_scores_list = []  # image level prediction
+        anomaly_maps = []  # pixel level prediction
+        gt_list = []  # image level ground truth
+        img_masks = []  # pixel level ground truth
+
+        image_path_list = []
+        test_dataset = self.load_datasets(category)
+        test_dataloader = torch.utils.data.DataLoader(
+            test_dataset,
+            batch_size=1,
+            shuffle=False,
+            num_workers=0,
+            pin_memory=True,
+        )
+
+        for image_info in tqdm(test_dataloader):
+            if not isinstance(image_info, dict):
+                raise ValueError("Encountered non-dict image in dataloader")
+
+            del image_info["image"]
+
+            image_path = image_info["image_path"][0]
+            image_path_list.extend(image_path)
+
+            img_masks.append(image_info["mask"])
+            gt_list.extend(list(image_info["is_anomaly"].numpy()))
+
+            image_scores = self.get_scores_for_image(image_path)
+            cls_scores = image_scores["img_level_score"]
+            anomaly_maps_iter = image_scores["pix_level_score"]
+
+            cls_scores_list.append(cls_scores)
+            anomaly_maps.append(anomaly_maps_iter)
+
+        pr_sp = np.array(cls_scores_list)
+        gt_sp = np.array(gt_list)
+        pr_px = np.array(anomaly_maps)
+        gt_px = torch.cat(img_masks, dim=0).numpy().astype(np.int32)
+
+        assert pr_px.shape[1:] == (
+            1,
+            224,
+            224,
+        ), "Predicted output scores do not meet the expected shape!"
+        assert gt_px.shape[1:] == (
+            1,
+            224,
+            224,
+        ), "Loaded ground truth maps do not meet the expected shape!"
+
+        return gt_sp, pr_sp, gt_px, pr_px, image_path_list

evaluation/utils/json_helpers.py

@@ -0,0 +1,46 @@
+# -----------------------------------------------------------------------------
+#  Do Not Alter This File!
+# -----------------------------------------------------------------------------
+#  The following code is part of the logic used for loading and evaluating your
+#  output scores. Please DO NOT modify this section, as upon your submission,
+#  the whole evaluation logic will be overwritten by the original code.
+# -----------------------------------------------------------------------------
+
+import json
+import numpy as np
+
+
+class NumpyEncoder(json.JSONEncoder):
+    """Special json encoder for numpy types"""
+
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.ndarray):
+            return {
+                "__ndarray__": obj.tolist(),
+                "dtype": str(obj.dtype),
+                "shape": obj.shape,
+            }
+        else:
+            return super(NumpyEncoder, self).default(obj)
+
+
+def dict_to_json(dct, filename):
+    """Save a dictionary to a JSON file"""
+    with open(filename, "w") as f:
+        json.dump(dct, f, cls=NumpyEncoder)
+
+
+def json_to_dict(filename):
+    """Load a JSON file and convert it back to a dictionary of NumPy arrays"""
+    with open(filename, "r") as f:
+        dct = json.load(f)
+
+    for k, v in dct.items():
+        if isinstance(v, dict) and "__ndarray__" in v:
+            dct[k] = np.array(v["__ndarray__"], dtype=v["dtype"]).reshape(v["shape"])
+
+    return dct

evaluation/utils/metrics.py

@@ -0,0 +1,78 @@
+# -----------------------------------------------------------------------------
+#  Do Not Alter This File!
+# -----------------------------------------------------------------------------
+#  The following code is part of the logic used for loading and evaluating your
+#  output scores. Please DO NOT modify this section, as upon your submission,
+#  the whole evaluation logic will be overwritten by the original code.
+# -----------------------------------------------------------------------------
+
+import numpy as np
+from sklearn.metrics import (
+    auc,
+    roc_auc_score,
+    average_precision_score,
+    precision_recall_curve,
+)
+from skimage import measure
+
+
+# ref: https://github.com/gudovskiy/cflow-ad/blob/master/train.py
+def cal_pro_score(masks, amaps, max_step=200, expect_fpr=0.3):
+    binary_amaps = np.zeros_like(amaps, dtype=bool)
+    min_th, max_th = amaps.min(), amaps.max()
+    delta = (max_th - min_th) / max_step
+    pros, fprs, ths = [], [], []
+    for th in np.arange(min_th, max_th, delta):
+        binary_amaps[amaps <= th], binary_amaps[amaps > th] = 0, 1
+        pro = []
+        for binary_amap, mask in zip(binary_amaps, masks):
+            for region in measure.regionprops(measure.label(mask)):
+                tp_pixels = binary_amap[region.coords[:, 0], region.coords[:, 1]].sum()
+                pro.append(tp_pixels / region.area)
+        inverse_masks = 1 - masks
+        fp_pixels = np.logical_and(inverse_masks, binary_amaps).sum()
+        fpr = fp_pixels / inverse_masks.sum()
+        pros.append(np.array(pro).mean())
+        fprs.append(fpr)
+        ths.append(th)
+    pros, fprs, ths = np.array(pros), np.array(fprs), np.array(ths)
+    idxes = fprs < expect_fpr
+    fprs = fprs[idxes]
+    fprs = (fprs - fprs.min()) / (fprs.max() - fprs.min())
+    pro_auc = auc(fprs, pros[idxes])
+    return pro_auc
+
+
+def compute_metrics(gt_sp=None, pr_sp=None, gt_px=None, pr_px=None):
+    # classification
+    if (
+        gt_sp is None
+        or pr_sp is None
+        or gt_sp.sum() == 0
+        or gt_sp.sum() == gt_sp.shape[0]
+    ):
+        auroc_sp, f1_sp, ap_sp = 0, 0, 0
+    else:
+        auroc_sp = roc_auc_score(gt_sp, pr_sp)
+        ap_sp = average_precision_score(gt_sp, pr_sp)
+        precisions, recalls, thresholds = precision_recall_curve(gt_sp, pr_sp)
+        f1_scores = (2 * precisions * recalls) / (precisions + recalls)
+        f1_sp = np.max(f1_scores[np.isfinite(f1_scores)])
+
+    # segmentation
+    if gt_px is None or pr_px is None or gt_px.sum() == 0:
+        auroc_px, f1_px, ap_px, aupro = 0, 0, 0, 0
+    else:
+        auroc_px = roc_auc_score(gt_px.ravel(), pr_px.ravel())
+        ap_px = average_precision_score(gt_px.ravel(), pr_px.ravel())
+        precisions, recalls, thresholds = precision_recall_curve(
+            gt_px.ravel(), pr_px.ravel()
+        )
+        f1_scores = (2 * precisions * recalls) / (precisions + recalls)
+        f1_px = np.max(f1_scores[np.isfinite(f1_scores)])
+        aupro = cal_pro_score(gt_px.squeeze(), pr_px.squeeze())
+
+    image_metric = [auroc_sp, f1_sp, ap_sp]
+    pixel_metric = [auroc_px, f1_px, ap_px, aupro]
+
+    return image_metric, pixel_metric

install.sh

@@ -0,0 +1,15 @@
+# add dependencies
+# python395_cuda113_pytorch1101
+# please change dataset root in ./config.py according to your specifications
+
+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
+
+
+
+
+
+

loss.py

@@ -0,0 +1,189 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from math import exp
+
+class FocalLoss(nn.Module):
+    """
+    copy from: https://github.com/Hsuxu/Loss_ToolBox-PyTorch/blob/master/FocalLoss/FocalLoss.py
+    This is a implementation of Focal Loss with smooth label cross entropy supported which is proposed in
+    'Focal Loss for Dense Object Detection. (https://arxiv.org/abs/1708.02002)'
+        Focal_Loss= -1*alpha*(1-pt)*log(pt)
+    :param alpha: (tensor) 3D or 4D the scalar factor for this criterion
+    :param gamma: (float,double) gamma > 0 reduces the relative loss for well-classified examples (p>0.5) putting more
+                    focus on hard misclassified example
+    :param smooth: (float,double) smooth value when cross entropy
+    :param balance_index: (int) balance class index, should be specific when alpha is float
+    :param size_average: (bool, optional) By default, the losses are averaged over each loss element in the batch.
+    """
+
+    def __init__(self, apply_nonlin=None, alpha=None, gamma=2, balance_index=0, smooth=1e-5, size_average=True):
+        super(FocalLoss, self).__init__()
+        self.apply_nonlin = apply_nonlin
+        self.alpha = alpha
+        self.gamma = gamma
+        self.balance_index = balance_index
+        self.smooth = smooth
+        self.size_average = size_average
+
+        if self.smooth is not None:
+            if self.smooth < 0 or self.smooth > 1.0:
+                raise ValueError('smooth value should be in [0,1]')
+
+    def forward(self, logit, target):
+        if self.apply_nonlin is not None:
+            logit = self.apply_nonlin(logit)
+        num_class = logit.shape[1]
+
+        if logit.dim() > 2:
+            # N,C,d1,d2 -> N,C,m (m=d1*d2*...)
+            logit = logit.view(logit.size(0), logit.size(1), -1)
+            logit = logit.permute(0, 2, 1).contiguous()
+            logit = logit.view(-1, logit.size(-1))
+        target = torch.squeeze(target, 1)
+        target = target.view(-1, 1)
+        alpha = self.alpha
+
+        if alpha is None:
+            alpha = torch.ones(num_class, 1)
+        elif isinstance(alpha, (list, np.ndarray)):
+            assert len(alpha) == num_class
+            alpha = torch.FloatTensor(alpha).view(num_class, 1)
+            alpha = alpha / alpha.sum()
+        elif isinstance(alpha, float):
+            alpha = torch.ones(num_class, 1)
+            alpha = alpha * (1 - self.alpha)
+            alpha[self.balance_index] = self.alpha
+
+        else:
+            raise TypeError('Not support alpha type')
+
+        if alpha.device != logit.device:
+            alpha = alpha.to(logit.device)
+
+        idx = target.cpu().long()
+
+        one_hot_key = torch.FloatTensor(target.size(0), num_class).zero_()
+        one_hot_key = one_hot_key.scatter_(1, idx, 1)
+        if one_hot_key.device != logit.device:
+            one_hot_key = one_hot_key.to(logit.device)
+
+        if self.smooth:
+            one_hot_key = torch.clamp(
+                one_hot_key, self.smooth / (num_class - 1), 1.0 - self.smooth)
+        pt = (one_hot_key * logit).sum(1) + self.smooth
+        logpt = pt.log()
+
+        gamma = self.gamma
+
+        alpha = alpha[idx]
+        alpha = torch.squeeze(alpha)
+        loss = -1 * alpha * torch.pow((1 - pt), gamma) * logpt
+
+        if self.size_average:
+            loss = loss.mean()
+        return loss
+
+
+class BinaryDiceLoss(nn.Module):
+    def __init__(self):
+        super(BinaryDiceLoss, self).__init__()
+
+    def forward(self, input, targets):
+        # 获取每个批次的大小 N
+        N = targets.size()[0]
+        # 平滑变量
+        smooth = 1
+        # 将宽高 reshape 到同一纬度
+        input_flat = input.view(N, -1)
+        targets_flat = targets.view(N, -1)
+
+        # 计算交集
+        intersection = input_flat * targets_flat
+        N_dice_eff = (2 * intersection.sum(1) + smooth) / (input_flat.sum(1) + targets_flat.sum(1) + smooth)
+        # 计算一个批次中平均每张图的损失
+        loss = 1 - N_dice_eff.sum() / N
+        return loss
+
+
+
+
+class ConADLoss(nn.Module):
+    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
+    It also supports the unsupervised contrastive loss in SimCLR"""
+    def __init__(self, contrast_mode='all',random_anchors=10):
+        super(ConADLoss, self).__init__()
+        assert contrast_mode in ['all', 'mean', 'random']
+        self.contrast_mode = contrast_mode
+        self.random_anchors = random_anchors
+    def forward(self, features, labels):
+        """Compute loss for model. If both `labels` and `mask` are None,
+        it degenerates to SimCLR unsupervised loss:
+        https://arxiv.org/pdf/2002.05709.pdf
+
+        Args:
+            features: hidden vector of shape [bsz, C, ...].
+            labels: ground truth of shape [bsz, 1, ...]., where 1 denotes to abnormal, and 0 denotes to normal
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if features.is_cuda
+                  else torch.device('cpu'))
+        if len(features.shape) != len(labels.shape):
+            raise ValueError('`features` needs to have the same dimensions with labels')
+
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, C, ...],'
+                             'at least 3 dimensions are required')
+
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+            labels = labels.view(labels.shape[0], labels.shape[1], -1)
+
+        labels = labels.squeeze()
+        batch_size = features.shape[0]
+
+        C = features.shape[1]
+        normal_feats = features[:, :, labels == 0]
+        abnormal_feats = features[:, :, labels == 1]
+
+        normal_feats = normal_feats.permute((1, 0, 2)).contiguous().view(C, -1)
+        abnormal_feats = abnormal_feats.permute((1, 0, 2)).contiguous().view(C, -1)
+
+        contrast_count = normal_feats.shape[1]
+        contrast_feature = normal_feats
+
+        if self.contrast_mode == 'mean':
+            anchor_feature = torch.mean(normal_feats, dim=1)
+            anchor_feature = F.normalize(anchor_feature, dim=0, p=2)
+            anchor_count = 1
+        elif self.contrast_mode == 'all':
+            anchor_feature = contrast_feature
+            anchor_count = contrast_count
+        elif self.contrast_mode == 'random':
+            dim_to_sample = 1
+            num_samples = min(self.random_anchors, contrast_count)
+            permuted_indices = torch.randperm(normal_feats.size(dim_to_sample)).to(normal_feats.device)
+            selected_indices = permuted_indices[:num_samples]
+            anchor_feature = normal_feats.index_select(dim_to_sample, selected_indices)
+        else:
+            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))
+
+        # compute logits
+        # maximize similarity
+        anchor_dot_normal = torch.matmul(anchor_feature.T, normal_feats).mean()
+
+        # minimize similarity
+        anchor_dot_abnormal = torch.matmul(anchor_feature.T, abnormal_feats).mean()
+
+        loss = 0
+        if normal_feats.shape[1] > 0:
+            loss -= anchor_dot_normal
+        if abnormal_feats.shape[1] > 0:
+            loss += anchor_dot_abnormal
+
+        loss = torch.exp(loss)
+
+        return loss