dataset.py

import os
import pandas as pd
import numpy as np
from enum import Enum

import PIL
import torch
from torchvision import transforms
from omegaconf import ListConfig

MVTecAD_CLASSNAMES = [
    "bottle",
    "cable",
    "capsule",
    "carpet",
    "grid",
    "hazelnut",
    "leather",
    "metal_nut",
    "pill",
    "screw",
    "tile",
    "toothbrush",
    "transistor",
    "wood",
    "zipper",
]

VisA_CLASSNAMES = [
    'candle',
    'capsules',
    'cashew',
    'chewinggum',
    'fryum',
    'macaroni1',
    'macaroni2',
    'pcb1',
    'pcb2',
    'pcb3',
    'pcb4',
    'pipe_fryum'
]

BTAD_CLASSNAMES = [
    '01',
    '02',
    '03'
]

MPDD_CLASSNAMES = [
    'bracket_black', 
    'bracket_brown', 
    'bracket_white', 
    'connector', 
    'metal_plate', 
    'tubes'
]

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 MVTecDataset(torch.utils.data.Dataset):
    """
    PyTorch Dataset for MVTec.
    """

    def __init__(
        self,
        source,
        classname,
        resize=256,
        imagesize=224,
        split=DatasetSplit.TRAIN,
        train_val_split=1.0,
        class_list=MVTecAD_CLASSNAMES,
        normal_state='good',
        **kwargs,
    ):
        """
        Args:
            source: [str]. Path to the MVTec data folder.
            classname: [str or None]. Name of MVTec class that should be
                       provided in this dataset. If None, the datasets
                       iterates over all available images.
            resize: [int]. (Square) Size the loaded image initially gets
                    resized to.
            imagesize: [int]. (Square) Size the resized loaded image gets
                       (center-)cropped to.
            split: [enum-option]. Indicates if training or test split of the
                   data should be used. Has to be an option taken from
                   DatasetSplit, e.g. mvtec.DatasetSplit.TRAIN. Note that
                   mvtec.DatasetSplit.TEST will also load mask data.
        """
        super().__init__()
        self.source = source
        self.split = split
        self.classname = classname
        self.train_val_split = train_val_split
        
        self.class_list = class_list
        self.normal_state = normal_state

        self.imgpaths_per_class, self.data_to_iterate = self.get_image_data()

        self.transform_img = [
            transforms.Resize(resize),
            transforms.CenterCrop(imagesize),
            transforms.ToTensor(),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ]
        self.transform_img = transforms.Compose(self.transform_img)

        self.transform_mask = [
            transforms.Resize(resize),
            transforms.CenterCrop(imagesize),
            transforms.ToTensor(),
        ]
        self.transform_mask = transforms.Compose(self.transform_mask)

        self.imagesize = (3, imagesize, imagesize)

    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)
            mask = self.transform_mask(mask)
        else:
            mask = torch.zeros([1, *image.size()[1:]])
            
        mask[mask > 0.0], mask[mask == 0.0] = 1, 0

        return {
            "image": image,
            "mask": mask,
            "classname": classname,
            "anomaly": anomaly,
            "is_anomaly": int(anomaly != self.normal_state),
            "image_name": "/".join(image_path.split("/")[-4:]),
            "image_path": image_path,
        }

    def __len__(self):
        return len(self.data_to_iterate)

    def get_image_data(self):
        if self.classname == 'all':
            self.classnames_to_use = self.class_list
        elif isinstance(self.classname, list) or isinstance(self.classname, ListConfig):
            for c in self.classname:
                assert c in self.class_list
            self.classnames_to_use = self.classname
        else:
            self.classnames_to_use = [self.classname]
        
        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.train_val_split < 1.0:
                    n_images = len(imgpaths_per_class[classname][anomaly])
                    train_val_split_idx = int(n_images * self.train_val_split)
                    if self.split == DatasetSplit.TRAIN:
                        imgpaths_per_class[classname][anomaly] = imgpaths_per_class[
                            classname
                        ][anomaly][:train_val_split_idx]
                    elif self.split == DatasetSplit.VAL:
                        imgpaths_per_class[classname][anomaly] = imgpaths_per_class[
                            classname
                        ][anomaly][train_val_split_idx:]

                if self.split == DatasetSplit.TEST and anomaly != self.normal_state:
                    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][self.normal_state] = None

        # Unrolls the data dictionary to an easy-to-iterate list.
        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 != self.normal_state:
                        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
    

class BTADDataset(MVTecDataset):
    """
    PyTorch Dataset for VisA.
    """

    def __init__(
        self,
        source,
        classname,
        resize=256,
        imagesize=224,
        split=DatasetSplit.TRAIN,
        train_val_split=1.0,
        class_list=BTAD_CLASSNAMES,
        normal_state='ok',
        **kwargs,
    ):
        """
        Args:
            source: [str]. Path to the MVTec data folder.
            classname: [str or None]. Name of MVTec class that should be
                       provided in this dataset. If None, the datasets
                       iterates over all available images.
            resize: [int]. (Square) Size the loaded image initially gets
                    resized to.
            imagesize: [int]. (Square) Size the resized loaded image gets
                       (center-)cropped to.
            split: [enum-option]. Indicates if training or test split of the
                   data should be used. Has to be an option taken from
                   DatasetSplit, e.g. mvtec.DatasetSplit.TRAIN. Note that
                   mvtec.DatasetSplit.TEST will also load mask data.
        """
        super().__init__(
            source = source,
            classname = classname,
            resize = resize,
            imagesize = imagesize,
            split = split,
            train_val_split = train_val_split,
            class_list = class_list,
            normal_state = normal_state
        )

class MPDDDataset(MVTecDataset):
    """
    PyTorch Dataset for VisA.
    """

    def __init__(
        self,
        source,
        classname,
        imagesize=224,
        split=DatasetSplit.TRAIN,
        train_val_split=1.0,
        class_list=MPDD_CLASSNAMES,
        normal_state='good',
        **kwargs,
    ):
        """
        Args:
            source: [str]. Path to the MVTec data folder.
            classname: [str or None]. Name of MVTec class that should be
                       provided in this dataset. If None, the datasets
                       iterates over all available images.
            resize: [int]. (Square) Size the loaded image initially gets
                    resized to.
            imagesize: [int]. (Square) Size the resized loaded image gets
                       (center-)cropped to.
            split: [enum-option]. Indicates if training or test split of the
                   data should be used. Has to be an option taken from
                   DatasetSplit, e.g. mvtec.DatasetSplit.TRAIN. Note that
                   mvtec.DatasetSplit.TEST will also load mask data.
        """
        super().__init__(
            source = source,
            classname = classname,
            imagesize = imagesize,
            split = split,
            train_val_split = train_val_split,
            class_list = class_list,
            normal_state = normal_state
        )
        

class VisADataset(MVTecDataset):
    """
    PyTorch Dataset for VisA.
    """

    def __init__(
        self,
        source,
        classname,
        resize=256,
        imagesize=224,
        split=DatasetSplit.TRAIN,
        train_val_split=1.0,
        class_list=VisA_CLASSNAMES,
        normal_state='good',
        **kwargs,
    ):
        """
        Args:
            source: [str]. Path to the MVTec data folder.
            classname: [str or None]. Name of MVTec class that should be
                       provided in this dataset. If None, the datasets
                       iterates over all available images.
            resize: [int]. (Square) Size the loaded image initially gets
                    resized to.
            imagesize: [int]. (Square) Size the resized loaded image gets
                       (center-)cropped to.
            split: [enum-option]. Indicates if training or test split of the
                   data should be used. Has to be an option taken from
                   DatasetSplit, e.g. mvtec.DatasetSplit.TRAIN. Note that
                   mvtec.DatasetSplit.TEST will also load mask data.
        """
        super().__init__(
            source = source,
            classname = classname,
            resize = resize,
            imagesize = imagesize,
            split = split,
            train_val_split = train_val_split,
            class_list = class_list,
            normal_state = normal_state
        )
        
    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)
            mask = np.array(mask)
            mask[mask > 0] = 255
            mask = PIL.Image.fromarray(mask)
            mask = self.transform_mask(mask)
        else:
            mask = torch.zeros([1, *image.size()[1:]])
            
        mask[mask > 0.0], mask[mask == 0.0] = 1, 0

        return {
            "image": image,
            "mask": mask,
            "classname": classname,
            "anomaly": anomaly,
            "is_anomaly": int(anomaly != self.normal_state),
            "image_name": "/".join(image_path.split("/")[-4:]),
            "image_path": image_path,
        }

    def get_image_data(self):
        self.classnames_to_use = [self.classname] if self.classname != 'all' else self.class_list
        
        split_csv = pd.read_csv(os.path.join(self.source, 'split_csv', '1cls.csv'))
        split_csv_cls = split_csv[
            split_csv['object'].isin(self.classnames_to_use) & (split_csv['split'] == self.split.value)
        ]
        split_csv_cls = split_csv_cls.drop('split', axis=1)
        
        # data_to_iterate
        split_csv_cls.loc[split_csv_cls['label'] == 'normal', 'label'] = self.normal_state
        split_csv_cls.loc[split_csv_cls['label'] == self.normal_state, 'mask'] = None
        split_csv_cls['image'] = split_csv_cls['image'].apply(lambda x: os.path.join(self.source, x))
        split_csv_cls['mask'] = split_csv_cls['mask'].apply(lambda x: os.path.join(self.source, x) if x != None else x)
        
        data_to_iterate = list(map(list, split_csv_cls.values))
        
        # imgpaths_per_class
        imgpaths_per_class = {}
        for cls in self.classnames_to_use:
            imgpaths_per_class[cls] = {}

            split_csv_cls_path = split_csv_cls[split_csv_cls['object']==cls]

            for label in split_csv_cls_path.label.unique():
                imgpaths_per_class[cls][label] = split_csv_cls_path.loc[split_csv_cls_path['label'] == label, 'image'].tolist()

        return imgpaths_per_class, data_to_iterate
    

def create_dataset(
    dataname: str, 
    source: str, 
    classname: str, 
    resize: int = 256, 
    imagesize: int = 224, 
    train_val_split: float = 1.0,
    split: DatasetSplit = DatasetSplit.TRAIN
):
    dataset = eval(f'{dataname}Dataset')(
        source          = source,
        classname       = classname,
        resize          = resize,
        imagesize       = imagesize,
        split           = split,
        train_val_split = train_val_split
    )
    
    return dataset