The initial code

Author: pjl1995

License_M3DM.txt

@@ -1,2 +1,114 @@
-# AnomalyDetection-M3DM
-Code for CVPR 2023 paper "Multimodal Industrial Anomaly Detection via Hybrid Fusion"
+# Multimodal Industrial Anomaly Detection via Hybrid Fusion
+
+![piplien](figures/pipeline.png)
+- `The pipeline of  Multi-3D-Memory (M3DM).` Our M3DM contains three important parts: (1) **Point Feature Alignment** (PFA) converts Point Group features to plane features with interpolation and project operation, $\text{FPS}$ is the farthest point sampling and $\mathcal F_{pt}$ is a pretrained Point Transformer; (2) **Unsupervised Feature Fusion** (UFF) fuses point feature and image feature together with a patch-wise contrastive loss $\mathcal L_{con}$, where $\mathcal F_{rgb}$ is a Vision Transformer, $\chi_{rgb},\chi_{pt}$ are MLP layers and $\sigma_r, \sigma_p$ are single fully connected layers; (3) **Decision Layer Fusion** (DLF) combines multimodal information with multiple memory banks and makes the final decision with 2 learnable modules $\mathcal D_a, \mathcal D_s$ for anomaly detection and segmentation, where $\mathcal{M}_{rgb}, \mathcal{M}_{fs}, \mathcal{M}_{pt}$ are memory banks, $\phi, \psi$ are score function for single memory bank detection and segmentation, and  $\mathcal{P}$ is the memory bank building algorithm.
+
+## Setup
+
+We implement this repo with the following environment:
+- Python 3.8
+- Pytorch 1.9.0
+- CUDA 11.3
+
+Install the other package via:
+
+``` bash
+pip install -r requirement.txt
+# install knn_cuda
+pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl
+# install pointnet2_ops_lib
+pip install "git+git://github.com/erikwijmans/Pointnet2_PyTorch.git#egg=pointnet2_ops&subdirectory=pointnet2_ops_lib"
+```
+
+## Data Download and Preprocess
+
+### Dataset
+
+- The `MVTec-3D AD` dataset can be download from the [Official Website of MVTec-3D AD](https://www.mvtec.com/company/research/datasets/mvtec-3d-ad). 
+
+- The `Eyecandies` dataset can be download from the [Official Website of Eyecandies](https://eyecan-ai.github.io/eyecandies/). 
+
+After download, put the dataset in `dataset` folder.
+
+### Datapreprocess
+
+
+To run the preprocessing 
+```bash
+python utils/preprocessing.py datasets/mvtec3d/
+```
+
+It may take a few hours to run the preprocessing. 
+
+### Checkpoints
+
+The following table lists the pretrain model used in M3DM:
+
+| Backbone          | Pretrain Method                                                                                                                                                                 |
+| ----------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| Point Transformer | [Point-MAE](https://github.com/Pang-Yatian/Point-MAE/releases/download/main/pretrain.pth)                                                                                       |
+| Point Transformer | [Point-Bert](https://cloud.tsinghua.edu.cn/f/202b29805eea45d7be92/?dl=1)                                                                                                        |
+| ViT-b/8           | [DINO](https://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth)                                                                                   |
+| ViT-b/8           | [Supervised ImageNet 1K](https://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_224.npz) |
+| ViT-b/8           | [Supervised ImageNet 21K](https://storage.googleapis.com/vit_models/augreg/B_8-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz)                                        |
+| ViT-s/8           | [DINO](https://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth)                                                                               |
+
+Put the checkpoint files in `checkpoints` folder.
+
+## Train and Test
+
+Train and test the double lib version and save the feature for UFF training:
+
+```bash
+python3 main.py \
+--method_name DINO+Point_MAE \
+--memory_bank multiple \
+--rgb_backbone_name vit_base_patch8_224_dino \
+--xyz_backbone_name Point_MAE \
+--save_feature True \
+```
+
+Train the UFF:
+
+```bash
+OMP_NUM_THREADS=1 python3 -m torch.distributed.launch --nproc_per_node=1 fusion_pretrain.py    \
+--accum_iter 16 \
+--lr 0.003 \
+--batch_size 16 \
+--data_path datasets/patch_lib \
+--output_dir checkpoints \
+```
+
+Train and test the full setting with the following command:
+
+```bash
+python3 main.py \
+--method_name DINO+Point_MAE+Fusion \
+--use_uff \
+--memory_bank multiple \
+--rgb_backbone_name vit_base_patch8_224_dino \
+--xyz_backbone_name Point_MAE \
+--fusion_module_path checkpoints/{FUSION_CHECKPOINT}.pth \
+```
+
+Note: if you set `--method_name DINO` or `--method_name Point_MAE`, set `--memory_bank single` at the same time. 
+
+
+
+If you find this repository useful for your research, please use the following.
+
+```bibtex
+
+@misc{wang2023multimodal,
+  title={Multimodal Industrial Anomaly Detection via Hybrid Fusion},
+  author={Wang, Yue and Peng, Jinlong and Zhang, Jiangning and Yi, Ran and Wang, Yabiao and Wang, Chengjie},
+  year={2023},
+  eprint={2303.00601},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV}
+}
+```
+
+## Thanks
+
+Our repo is built on [3D-ADS](https://github.com/eliahuhorwitz/3D-ADS) and [MoCo-v3](https://github.com/facebookresearch/moco-v3), thanks their extraordinary works!

README.md

@@ -0,0 +1,188 @@
+import os
+from PIL import Image
+from torchvision import transforms
+import glob
+from torch.utils.data import Dataset
+from utils.mvtec3d_util import *
+from torch.utils.data import DataLoader
+import numpy as np
+
+def eyecandies_classes():
+    return [
+        'CandyCane',
+        'ChocolateCookie',
+        'ChocolatePraline',
+        'Confetto',
+        'GummyBear',
+        'HazelnutTruffle',
+        'LicoriceSandwich',
+        'Lollipop',
+        'Marshmallow',
+        'PeppermintCandy',   
+    ]
+
+def mvtec3d_classes():
+    return [
+        "bagel",
+        "cable_gland",
+        "carrot",
+        "cookie",
+        "dowel",
+        "foam",
+        "peach",
+        "potato",
+        "rope",
+        "tire",
+    ]
+
+RGB_SIZE = 224
+
+class BaseAnomalyDetectionDataset(Dataset):
+
+    def __init__(self, split, class_name, img_size, dataset_path='datasets/eyecandies_preprocessed'):
+        self.IMAGENET_MEAN = [0.485, 0.456, 0.406]
+        self.IMAGENET_STD = [0.229, 0.224, 0.225]
+        self.cls = class_name
+        self.size = img_size
+        self.img_path = os.path.join(dataset_path, self.cls, split)
+        self.rgb_transform = transforms.Compose(
+            [transforms.Resize((RGB_SIZE, RGB_SIZE), interpolation=transforms.InterpolationMode.BICUBIC),
+             transforms.ToTensor(),
+             transforms.Normalize(mean=self.IMAGENET_MEAN, std=self.IMAGENET_STD)])
+
+class PreTrainTensorDataset(Dataset):
+    def __init__(self, root_path):
+        super().__init__()
+        self.root_path = root_path
+        self.tensor_paths = os.listdir(self.root_path)
+
+
+    def __len__(self):
+        return len(self.tensor_paths)
+
+    def __getitem__(self, idx):
+        tensor_path = self.tensor_paths[idx]
+
+        tensor = torch.load(os.path.join(self.root_path, tensor_path))
+
+        label = 0
+
+        return tensor, label
+
+class TrainDataset(BaseAnomalyDetectionDataset):
+    def __init__(self, class_name, img_size, dataset_path='datasets/eyecandies_preprocessed'):
+        super().__init__(split="train", class_name=class_name, img_size=img_size, dataset_path=dataset_path)
+        self.img_paths, self.labels = self.load_dataset()  # self.labels => good : 0, anomaly : 1
+
+    def load_dataset(self):
+        img_tot_paths = []
+        tot_labels = []
+        rgb_paths = glob.glob(os.path.join(self.img_path, 'good', 'rgb') + "/*.png")
+        tiff_paths = glob.glob(os.path.join(self.img_path, 'good', 'xyz') + "/*.tiff")
+        rgb_paths.sort()
+        tiff_paths.sort()
+        sample_paths = list(zip(rgb_paths, tiff_paths))
+        img_tot_paths.extend(sample_paths)
+        tot_labels.extend([0] * len(sample_paths))
+        return img_tot_paths, tot_labels
+
+    def __len__(self):
+        return len(self.img_paths)
+
+    def __getitem__(self, idx):
+        img_path, label = self.img_paths[idx], self.labels[idx]
+        rgb_path = img_path[0]
+        tiff_path = img_path[1]
+        img = Image.open(rgb_path).convert('RGB')
+
+        img = self.rgb_transform(img)
+        organized_pc = read_tiff_organized_pc(tiff_path)
+        
+        depth_map_3channel = np.repeat(organized_pc_to_depth_map(organized_pc)[:, :, np.newaxis], 3, axis=2)
+        resized_depth_map_3channel = resize_organized_pc(depth_map_3channel)
+        resized_organized_pc = resize_organized_pc(organized_pc, target_height=self.size, target_width=self.size)
+        resized_organized_pc = resized_organized_pc.clone().detach().float()
+
+        return (img, resized_organized_pc, resized_depth_map_3channel), label
+
+
+class TestDataset(BaseAnomalyDetectionDataset):
+    def __init__(self, class_name, img_size, dataset_path='datasets/eyecandies_preprocessed'):
+        super().__init__(split="test", class_name=class_name, img_size=img_size, dataset_path=dataset_path)
+        self.gt_transform = transforms.Compose([
+            transforms.Resize((RGB_SIZE, RGB_SIZE), interpolation=transforms.InterpolationMode.NEAREST),
+            transforms.ToTensor()])
+        self.img_paths, self.gt_paths, self.labels = self.load_dataset()  # self.labels => good : 0, anomaly : 1
+
+    def load_dataset(self):
+        img_tot_paths = []
+        gt_tot_paths = []
+        tot_labels = []
+        defect_types = os.listdir(self.img_path)
+
+        for defect_type in defect_types:
+            if defect_type == 'good':
+                rgb_paths = glob.glob(os.path.join(self.img_path, defect_type, 'rgb') + "/*.png")
+                tiff_paths = glob.glob(os.path.join(self.img_path, defect_type, 'xyz') + "/*.tiff")
+                rgb_paths.sort()
+                tiff_paths.sort()
+                sample_paths = list(zip(rgb_paths, tiff_paths))
+                img_tot_paths.extend(sample_paths)
+                gt_tot_paths.extend([0] * len(sample_paths))
+                tot_labels.extend([0] * len(sample_paths))
+            else:
+                rgb_paths = glob.glob(os.path.join(self.img_path, defect_type, 'rgb') + "/*.png")
+                tiff_paths = glob.glob(os.path.join(self.img_path, defect_type, 'xyz') + "/*.tiff")
+                gt_paths = glob.glob(os.path.join(self.img_path, defect_type, 'gt') + "/*.png")
+                rgb_paths.sort()
+                tiff_paths.sort()
+                gt_paths.sort()
+                sample_paths = list(zip(rgb_paths, tiff_paths))
+
+                img_tot_paths.extend(sample_paths)
+                gt_tot_paths.extend(gt_paths)
+                tot_labels.extend([1] * len(sample_paths))
+
+        assert len(img_tot_paths) == len(gt_tot_paths), "Something wrong with test and ground truth pair!"
+
+        return img_tot_paths, gt_tot_paths, tot_labels
+
+    def __len__(self):
+        return len(self.img_paths)
+
+    def __getitem__(self, idx):
+        img_path, gt, label = self.img_paths[idx], self.gt_paths[idx], self.labels[idx]
+        rgb_path = img_path[0]
+        tiff_path = img_path[1]
+        img_original = Image.open(rgb_path).convert('RGB')
+        img = self.rgb_transform(img_original)
+
+        organized_pc = read_tiff_organized_pc(tiff_path)
+        depth_map_3channel = np.repeat(organized_pc_to_depth_map(organized_pc)[:, :, np.newaxis], 3, axis=2)
+        resized_depth_map_3channel = resize_organized_pc(depth_map_3channel)
+        resized_organized_pc = resize_organized_pc(organized_pc, target_height=self.size, target_width=self.size)
+        resized_organized_pc = resized_organized_pc.clone().detach().float()
+        
+
+        
+
+        if gt == 0:
+            gt = torch.zeros(
+                [1, resized_depth_map_3channel.size()[-2], resized_depth_map_3channel.size()[-2]])
+        else:
+            gt = Image.open(gt).convert('L')
+            gt = self.gt_transform(gt)
+            gt = torch.where(gt > 0.5, 1., .0)
+
+        return (img, resized_organized_pc, resized_depth_map_3channel), gt[:1], label, rgb_path
+
+
+def get_data_loader(split, class_name, img_size, args):
+    if split in ['train']:
+        dataset = TrainDataset(class_name=class_name, img_size=img_size, dataset_path=args.dataset_path)
+    elif split in ['test']:
+        dataset = TestDataset(class_name=class_name, img_size=img_size, dataset_path=args.dataset_path)
+
+    data_loader = DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=1, drop_last=False,
+                             pin_memory=True)
+    return data_loader

dataset.py

@@ -0,0 +1,75 @@
+import math
+import sys
+from typing import Iterable
+
+import torch
+
+import utils.misc as misc
+import utils.lr_sched as lr_sched
+
+
+def train_one_epoch(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    log_writer=None,
+                    args=None):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    for data_iter_step, (samples, _) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        
+        xyz_samples = samples[:,:,:1152].to(device, non_blocking=True)
+        rgb_samples = samples[:,:,1152:].to(device, non_blocking=True)
+
+        with torch.cuda.amp.autocast():
+            loss = model(xyz_samples, rgb_samples)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
+            """ We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar('train_loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}