quadra 0.0.1__py3-none-any.whl → 2.1.13__py3-none-any.whl

quadra/datamodules/ssl.py

@@ -0,0 +1,140 @@
+# pylint: disable=unsubscriptable-object
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+
+from quadra.datamodules.classification import ClassificationDataModule
+from quadra.datasets import TwoAugmentationDataset, TwoSetAugmentationDataset
+from quadra.utils import utils
+
+log = utils.get_logger(__name__)
+
+
+class SSLDataModule(ClassificationDataModule):
+    """Base class for all data modules for self supervised learning data modules.
+
+    Args:
+        data_path: Path to the data main folder.
+        augmentation_dataset: Augmentation dataset
+            for training dataset.
+        name: The name for the data module. Defaults to  "ssl_datamodule".
+        split_validation: Whether to split the validation set if . Defaults to True.
+        **kwargs: The keyword arguments for the classification data module. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        augmentation_dataset: TwoAugmentationDataset | TwoSetAugmentationDataset,
+        name: str = "ssl_datamodule",
+        split_validation: bool = True,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            data_path=data_path,
+            name=name,
+            **kwargs,
+        )
+        self.augmentation_dataset = augmentation_dataset
+        self.classifier_train_dataset: torch.utils.data.Dataset | None = None
+        self.split_validation = split_validation
+
+    def setup(self, stage: str | None = None) -> None:
+        """Setup data module based on stages of training."""
+        if stage == "fit":
+            self.train_dataset = self.dataset(
+                samples=self.train_data["samples"].tolist(),
+                targets=self.train_data["targets"].tolist(),
+                transform=self.train_transform,
+            )
+
+            if np.unique(self.train_data["targets"]).shape[0] > 1 and not self.split_validation:
+                self.classifier_train_dataset = self.dataset(
+                    samples=self.train_data["samples"].tolist(),
+                    targets=self.train_data["targets"].tolist(),
+                    transform=self.val_transform,
+                )
+                self.val_dataset = self.dataset(
+                    samples=self.val_data["samples"].tolist(),
+                    targets=self.val_data["targets"].tolist(),
+                    transform=self.val_transform,
+                )
+            else:
+                train_classifier_samples, val_samples, train_classifier_targets, val_targets = train_test_split(
+                    self.val_data["samples"],
+                    self.val_data["targets"],
+                    test_size=0.3,
+                    random_state=self.seed,
+                    stratify=self.val_data["targets"],
+                )
+
+                self.classifier_train_dataset = self.dataset(
+                    samples=train_classifier_samples,
+                    targets=train_classifier_targets,
+                    transform=self.test_transform,
+                )
+
+                self.val_dataset = self.dataset(
+                    samples=val_samples,
+                    targets=val_targets,
+                    transform=self.val_transform,
+                )
+
+                log.warning(
+                    "The training set contains only one class and cannot be used to train a classifier. To overcome "
+                    "this issue 70% of the validation set is used to train the classifier. The remaining will be used "
+                    "as standard validation. To disable this behaviour set the `split_validation` parameter to False."
+                )
+                self._check_train_dataset_config()
+        if stage == "test":
+            self.test_dataset = self.dataset(
+                samples=self.test_data["samples"].tolist(),
+                targets=self.test_data["targets"].tolist(),
+                transform=self.test_transform,
+            )
+
+    def _check_train_dataset_config(self):
+        """Check if train dataset is configured correctly."""
+        if self.train_dataset is None:
+            raise ValueError("Train dataset is not initialized")
+        if self.augmentation_dataset is None:
+            raise ValueError("Augmentation dataset is not initialized")
+        if self.train_dataset.transform is not None:
+            log.warning("Train dataset transform is not None. It will be applied before SSL augmentations")
+
+    def train_dataloader(self) -> DataLoader:
+        """Returns train dataloader."""
+        if not isinstance(self.train_dataset, torch.utils.data.Dataset):
+            raise ValueError("Train dataset is not a subclass of `torch.utils.data.Dataset`")
+        self.augmentation_dataset.dataset = self.train_dataset
+        loader = DataLoader(
+            self.augmentation_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+        )
+        return loader
+
+    def classifier_train_dataloader(self) -> DataLoader:
+        """Returns classifier train dataloader."""
+        if self.classifier_train_dataset is None:
+            raise ValueError("Classifier train dataset is not initialized")
+
+        loader = DataLoader(
+            self.classifier_train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+        )
+        return loader

quadra/datasets/__init__.py

@@ -0,0 +1,17 @@
+from .anomaly import AnomalyDataset
+from .classification import ClassificationDataset, ImageClassificationListDataset, MultilabelClassificationDataset
+from .patch import PatchSklearnClassificationTrainDataset
+from .segmentation import SegmentationDataset, SegmentationDatasetMulticlass
+from .ssl import TwoAugmentationDataset, TwoSetAugmentationDataset
+
+__all__ = [
+    "ImageClassificationListDataset",
+    "ClassificationDataset",
+    "SegmentationDataset",
+    "SegmentationDatasetMulticlass",
+    "PatchSklearnClassificationTrainDataset",
+    "MultilabelClassificationDataset",
+    "AnomalyDataset",
+    "TwoAugmentationDataset",
+    "TwoSetAugmentationDataset",
+]

quadra/datasets/anomaly.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+import os
+import random
+from pathlib import Path
+
+import albumentations as alb
+import cv2
+import numpy as np
+import pandas as pd
+from pandas import DataFrame
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from quadra.utils.utils import IMAGE_EXTENSIONS
+
+
+def create_validation_set_from_test_set(samples: DataFrame, seed: int = 0) -> DataFrame:
+    """Craete Validation Set from Test Set.
+
+    This function creates a validation set from test set by splitting both
+    normal and abnormal samples to two.
+
+    Args:
+        samples: Dataframe containing dataset info such as filenames, splits etc.
+        seed: Random seed to ensure reproducibility. Defaults to 0.
+    """
+    if seed > 0:
+        random.seed(seed)
+
+    # Split normal images.
+    normal_test_image_indices = samples.index[(samples.split == "test") & (samples.targets == "good")].to_list()
+    num_normal_valid_images = len(normal_test_image_indices) // 2
+
+    indices_to_sample = random.sample(population=normal_test_image_indices, k=num_normal_valid_images)
+    samples.loc[indices_to_sample, "split"] = "val"
+
+    # Split abnormal images.
+    abnormal_test_image_indices = samples.index[(samples.split == "test") & (samples.targets != "good")].to_list()
+    num_abnormal_valid_images = len(abnormal_test_image_indices) // 2
+
+    indices_to_sample = random.sample(population=abnormal_test_image_indices, k=num_abnormal_valid_images)
+    samples.loc[indices_to_sample, "split"] = "val"
+
+    return samples
+
+
+def split_normal_images_in_train_set(samples: DataFrame, split_ratio: float = 0.1, seed: int = 0) -> DataFrame:
+    """Split normal images in train set.
+
+        This function splits the normal images in training set and assigns the
+        values to the test set. This is particularly useful especially when the
+        test set does not contain any normal images.
+
+        This is important because when the test set doesn't have any normal images,
+        AUC computation fails due to having single class.
+
+    Args:
+        samples: Dataframe containing dataset info such as filenames, splits etc.
+        split_ratio: Train-Test normal image split ratio. Defaults to 0.1.
+        seed: Random seed to ensure reproducibility. Defaults to 0.
+
+    Returns:
+        Output dataframe where the part of the training set is assigned to test set.
+    """
+    if seed > 0:
+        random.seed(seed)
+
+    normal_train_image_indices = samples.index[(samples.split == "train") & (samples.targets == "good")].to_list()
+    num_normal_train_images = len(normal_train_image_indices)
+    num_normal_valid_images = int(num_normal_train_images * split_ratio)
+
+    indices_to_split_from_train_set = random.sample(population=normal_train_image_indices, k=num_normal_valid_images)
+    samples.loc[indices_to_split_from_train_set, "split"] = "test"
+
+    return samples
+
+
+def make_anomaly_dataset(
+    path: Path,
+    split: str | None = None,
+    split_ratio: float = 0.1,
+    seed: int = 0,
+    mask_suffix: str | None = None,
+    create_test_set_if_empty: bool = True,
+) -> DataFrame:
+    """Create dataframe by parsing a folder following the MVTec data file structure.
+
+    The files are expected to follow the structure:
+        path/to/dataset/split/label/image_filename.xyz
+        path/to/dataset/ground_truth/label/mask_filename.png
+
+    Masks MUST be png images, no other format is allowed
+    Split can be either train/val/test
+
+    This function creates a dataframe to store the parsed information based on the following format:
+    |---|---------------|-------|---------|--------------|-----------------------------------------------|-------------|
+    |   | path          | split | targets | samples      | mask_path                                     | label_index |
+    |---|---------------|-------|---------|--------------|-----------------------------------------------|-------------|
+    | 0 | datasets/name |  test |  defect | filename.xyz | ground_truth/defect/filename{mask_suffix}.png | 1           |
+    |---|---------------|-------|---------|--------------|-----------------------------------------------|-------------|
+
+    Args:
+        path: Path to dataset
+        split: Dataset split (i.e., either train or test). Defaults to None.
+        split_ratio: Ratio to split normal training images and add to the
+            test set in case test set doesn't contain any normal images.
+            Defaults to 0.1.
+        seed: Random seed to ensure reproducibility when splitting. Defaults to 0.
+        mask_suffix: String to append to the base filename to get the mask name, by default for MVTec dataset masks
+            are saved as imagename_mask.png in this case the parameter shoul be filled with "_mask"
+        create_test_set_if_empty: If True, create a test set if the test set is empty.
+
+
+    Example:
+        The following example shows how to get training samples from MVTec bottle category:
+
+        >>> root = Path('./MVTec')
+        >>> category = 'bottle'
+        >>> path = root / category
+        >>> path
+        PosixPath('MVTec/bottle')
+
+        >>> samples = make_anomaly_dataset(path, split='train', split_ratio=0.1, seed=0)
+        >>> samples.head()
+           path         split label image_path                           mask_path                   label_index
+        0  MVTec/bottle train good MVTec/bottle/train/good/105.png MVTec/bottle/ground_truth/good/105_mask.png 0
+        1  MVTec/bottle train good MVTec/bottle/train/good/017.png MVTec/bottle/ground_truth/good/017_mask.png 0
+        2  MVTec/bottle train good MVTec/bottle/train/good/137.png MVTec/bottle/ground_truth/good/137_mask.png 0
+        3  MVTec/bottle train good MVTec/bottle/train/good/152.png MVTec/bottle/ground_truth/good/152_mask.png 0
+        4  MVTec/bottle train good MVTec/bottle/train/good/109.png MVTec/bottle/ground_truth/good/109_mask.png 0
+
+    Returns:
+        An output dataframe containing samples for the requested split (ie., train or test)
+    """
+    samples_list = [
+        (str(path),) + filename.parts[-3:]
+        for filename in path.glob("**/*")
+        if filename.is_file()
+        and os.path.splitext(filename)[-1].lower() in IMAGE_EXTENSIONS
+        and ".ipynb_checkpoints" not in str(filename)
+    ]
+
+    if len(samples_list) == 0:
+        raise RuntimeError(f"Found 0 images in {path}")
+
+    samples_list.sort()
+
+    data = pd.DataFrame(samples_list, columns=["path", "split", "targets", "samples"])
+    data = data[data.split != "ground_truth"]
+
+    # Create mask_path column, masks MUST have png extension
+    data["mask_path"] = (
+        data.path
+        + "/ground_truth/"
+        + data.targets
+        + "/"
+        + data.samples.apply(lambda x: os.path.splitext(os.path.basename(x))[0])
+        + (f"{mask_suffix}.png" if mask_suffix is not None else ".png")
+    )
+
+    # Modify image_path column by converting to absolute path
+    data["samples"] = data.path + "/" + data.split + "/" + data.targets + "/" + data.samples
+
+    # Split the normal images in training set if test set doesn't
+    # contain any normal images. This is needed because AUC score
+    # cannot be computed based on 1-class
+    if sum((data.split == "test") & (data.targets == "good")) == 0 and create_test_set_if_empty:
+        data = split_normal_images_in_train_set(data, split_ratio, seed)
+
+    # Good images don't have mask
+    data.loc[(data.split == "test") & (data.targets == "good"), "mask_path"] = ""
+
+    # Create label index for normal (0), anomalous (1) and unknown (-1) images.
+    data.loc[data.targets == "good", "label_index"] = 0
+    data.loc[~data.targets.isin(["good", "unknown"]), "label_index"] = 1
+    data.loc[data.targets == "unknown", "label_index"] = -1
+    data.label_index = data.label_index.astype(int)
+
+    # Get the data frame for the split.
+    if split is not None and split in ["train", "val", "test"]:
+        data = data[data.split == split]
+        data = data.reset_index(drop=True)
+
+    return data
+
+
+class AnomalyDataset(Dataset):
+    """Anomaly Dataset.
+
+    Args:
+        transform: Albumentations compose.
+        task: ``classification`` or ``segmentation``
+        samples: Pandas dataframe containing samples following the same structure created by make_anomaly_dataset
+        valid_area_mask: Optional path to the mask to use to filter out the valid area of the image. If None, the
+            whole image is considered valid.
+        crop_area: Optional tuple of 4 integers (x1, y1, x2, y2) to crop the image to the specified area. If None, the
+            whole image is considered valid.
+    """
+
+    def __init__(
+        self,
+        transform: alb.Compose,
+        samples: DataFrame,
+        task: str = "segmentation",
+        valid_area_mask: str | None = None,
+        crop_area: tuple[int, int, int, int] | None = None,
+    ) -> None:
+        self.task = task
+        self.transform = transform
+
+        self.samples = samples
+        self.samples = self.samples.reset_index(drop=True)
+        self.split = self.samples.split.unique()[0]
+
+        self.crop_area = crop_area
+        self.valid_area_mask: np.ndarray | None = None
+
+        if valid_area_mask is not None:
+            if not os.path.exists(valid_area_mask):
+                raise RuntimeError(f"Valid area mask {valid_area_mask} does not exist.")
+
+            self.valid_area_mask = cv2.imread(valid_area_mask, 0) > 0  # type: ignore[operator]
+
+    def __len__(self) -> int:
+        """Get length of the dataset."""
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> dict[str, str | Tensor]:
+        """Get dataset item for the index ``index``.
+
+        Args:
+            index: Index to get the item.
+
+        Returns:
+            Dict of image tensor during training.
+            Otherwise, Dict containing image path, target path, image tensor, label and transformed bounding box.
+        """
+        item: dict[str, str | Tensor] = {}
+
+        image_path = self.samples.samples.iloc[index]
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        original_image_shape = image.shape
+        if self.valid_area_mask is not None:
+            image = image * self.valid_area_mask[:, :, np.newaxis]
+
+        if self.crop_area is not None:
+            image = image[self.crop_area[1] : self.crop_area[3], self.crop_area[0] : self.crop_area[2]]
+
+        label_index = self.samples.label_index[index]
+
+        if self.split == "train":
+            pre_processed = self.transform(image=image)
+            item = {"image": pre_processed["image"], "label": label_index}
+        elif self.split in ["val", "test"]:
+            item["image_path"] = image_path
+            item["label"] = label_index
+
+            if self.task == "segmentation":
+                mask_path = self.samples.mask_path[index]
+
+                # If good images have no associated mask create an empty one
+                if label_index == 0:
+                    mask = np.zeros(shape=original_image_shape[:2])
+                elif os.path.isfile(mask_path):
+                    mask = cv2.imread(mask_path, flags=0) / 255.0  # type: ignore[operator]
+                else:
+                    # We need ones in the mask to compute correctly at least image level f1 score
+                    mask = np.ones(shape=original_image_shape[:2])
+
+                if self.valid_area_mask is not None:
+                    mask = mask * self.valid_area_mask
+
+                if self.crop_area is not None:
+                    mask = mask[self.crop_area[1] : self.crop_area[3], self.crop_area[0] : self.crop_area[2]]
+
+                pre_processed = self.transform(image=image, mask=mask)
+
+                item["mask_path"] = mask_path
+                item["mask"] = pre_processed["mask"]
+            else:
+                pre_processed = self.transform(image=image)
+
+            item["image"] = pre_processed["image"]
+        return item

quadra/datasets/classification.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import warnings
+from collections.abc import Callable
+
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from quadra.utils.imaging import crop_image, keep_aspect_ratio_resize
+
+
+class ImageClassificationListDataset(Dataset):
+    """Standard classification dataset.
+
+    Args:
+        samples: List of paths to images to be read
+        targets: List of labels, one for every image
+            in samples
+        class_to_idx: mapping from classes
+            to unique indexes.
+            Defaults to None.
+        resize: Integer specifying the size of
+            a first optional resize keeping the aspect ratio: the smaller side
+            of the image will be resized to `resize`, while the longer will be
+            resized keeping the aspect ratio.
+            Defaults to None.
+        roi: Optional ROI, with
+            (x_upper_left, y_upper_left, x_bottom_right, y_bottom_right).
+            Defaults to None.
+        transform: Optional Albumentations
+            transform.
+            Defaults to None.
+        rgb: if False, image will be converted in grayscale
+        channel: 1 or 3. If rgb is True, then channel will be set at 3.
+        allow_missing_label: If set to false warn the user if the dataset contains missing labels
+    """
+
+    def __init__(
+        self,
+        samples: list[str],
+        targets: list[str | int],
+        class_to_idx: dict | None = None,
+        resize: int | None = None,
+        roi: tuple[int, int, int, int] | None = None,
+        transform: Callable | None = None,
+        rgb: bool = True,
+        channel: int = 3,
+        allow_missing_label: bool | None = False,
+    ):
+        super().__init__()
+        assert len(samples) == len(
+            targets
+        ), f"Samples ({len(samples)}) and targets ({len(targets)}) must have the same length"
+        # Setting the ROI
+        self.roi = roi
+
+        # Keep-Aspect-Ratio resize
+        self.resize = resize
+
+        if not allow_missing_label and None in targets:
+            warnings.warn(
+                (
+                    "Dataset contains empty targets but allow_missing_label is set to False, "
+                    "be careful because None labels will not work inside Dataloaders"
+                ),
+                UserWarning,
+                stacklevel=2,
+            )
+
+        targets = [-1 if target is None else target for target in targets]
+        # Data
+        self.x = np.array(samples)
+        self.y = np.array(targets)
+
+        if class_to_idx is None:
+            unique_targets = np.unique(targets)
+            class_to_idx = {c: i for i, c in enumerate(unique_targets)}
+
+        self.class_to_idx = class_to_idx
+        self.idx_to_class = {v: k for k, v in class_to_idx.items()}
+        self.samples = [
+            (path, self.class_to_idx[self.y[i]] if (self.y[i] != -1 and self.y[i] != "-1") else -1)
+            for i, path in enumerate(self.x)
+        ]
+
+        self.rgb = rgb
+        self.channel = 3 if rgb else channel
+
+        self.transform = transform
+
+    def __getitem__(self, idx) -> tuple[np.ndarray, np.ndarray]:
+        path, y = self.samples[idx]
+
+        # Load image
+        x = cv2.imread(str(path))
+        if self.rgb:
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
+        else:
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2GRAY)
+            x = cv2.cvtColor(x, cv2.COLOR_GRAY2RGB)
+
+        if self.channel == 1:
+            x = x[:, :, 0]
+
+        # Crop with ROI
+        if self.roi:
+            x = crop_image(x, self.roi)
+
+        # Resize keeping aspect ratio
+        if self.resize:
+            x = keep_aspect_ratio_resize(x, self.resize)
+
+        if self.transform:
+            aug = self.transform(image=x)
+            x = aug["image"]
+
+        return x, y
+
+    def __len__(self):
+        return len(self.samples)
+
+
+class ClassificationDataset(ImageClassificationListDataset):
+    """Custom Classification Dataset.
+
+    Args:
+        samples: List of paths to images
+        targets: List of targets
+        class_to_idx: Defaults to None.
+        resize: Resize image to this size. Defaults to None.
+        roi: Region of interest. Defaults to None.
+        transform: transform function. Defaults to None.
+        rgb: Use RGB space
+        channel: Number of channels. Defaults to 3.
+        random_padding: Random padding. Defaults to False.
+        circular_crop: Circular crop. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        samples: list[str],
+        targets: list[str | int],
+        class_to_idx: dict | None = None,
+        resize: int | None = None,
+        roi: tuple[int, int, int, int] | None = None,
+        transform: Callable | None = None,
+        rgb: bool = True,
+        channel: int = 3,
+        random_padding: bool = False,
+        circular_crop: bool = False,
+    ):
+        super().__init__(samples, targets, class_to_idx, resize, roi, transform, rgb, channel)
+        if transform is None:
+            self.transform = None
+
+        self.random_padding = random_padding
+        self.circular_crop = circular_crop
+
+    def __getitem__(self, idx):
+        path, y = self.samples[idx]
+        path = str(path)
+
+        # Load image
+        x = cv2.imread(path)
+        if self.rgb:
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
+        else:
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2GRAY)
+            x = cv2.cvtColor(x, cv2.COLOR_GRAY2RGB)
+
+        if self.transform is not None:
+            aug = self.transform(image=x)
+            x = aug["image"]
+
+        if self.channel == 1:
+            x = x[:1]
+
+        return x, y
+
+
+class MultilabelClassificationDataset(torch.utils.data.Dataset):
+    """Custom MultilabelClassification Dataset.
+
+    Args:
+        samples: list of paths to images.
+        targets: array of multiple targets per sample. The array must be a one-hot enoding.
+            It must have a shape of (n_samples, n_targets).
+        class_to_idx: Defaults to None.
+        transform: transform function. Defaults to None.
+        rgb: Use RGB space
+    """
+
+    def __init__(
+        self,
+        samples: list[str],
+        targets: np.ndarray,
+        class_to_idx: dict | None = None,
+        transform: Callable | None = None,
+        rgb: bool = True,
+    ):
+        super().__init__()
+        assert len(samples) == len(
+            targets
+        ), f"Samples ({len(samples)}) and targets ({len(targets)}) must have the same length"
+
+        # Data
+        self.x = samples
+        self.y = targets
+
+        # Class to idx and the other way around
+        if class_to_idx is None:
+            unique_targets = targets.shape[1]
+            class_to_idx = {c: i for i, c in enumerate(range(unique_targets))}
+        self.class_to_idx = class_to_idx
+        self.idx_to_class = {v: k for k, v in class_to_idx.items()}
+        self.samples = list(zip(self.x, self.y))
+        self.rgb = rgb
+        self.transform = transform
+
+    def __getitem__(self, idx):
+        path, y = self.samples[idx]
+        path = str(path)
+
+        # Load image
+        x = cv2.imread(path)
+        if self.rgb:
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
+        else:
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2GRAY)
+            x = cv2.cvtColor(x, cv2.COLOR_GRAY2RGB)
+
+        if self.transform is not None:
+            aug = self.transform(image=x)
+            x = aug["image"]
+
+        return x, torch.from_numpy(y).float()
+
+    def __len__(self):
+        return len(self.samples)