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

quadra/datasets/patch.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import os
+import random
+from collections.abc import Callable
+
+import cv2
+import h5py
+import numpy as np
+from torch.utils.data import Dataset
+
+from quadra.utils import utils
+from quadra.utils.imaging import keep_aspect_ratio_resize
+from quadra.utils.patch.dataset import compute_safe_patch_range, trisample
+
+log = utils.get_logger(__name__)
+
+
+class PatchSklearnClassificationTrainDataset(Dataset):
+    """Dataset used for patch sampling, it expects samples to be paths to h5 files containing all the required
+    information for patch sampling from images.
+
+    Args:
+        data_path: base path to the dataset
+        samples: Paths to h5 files
+        targets: Labels associated with each sample
+        class_to_idx: Mapping between class and corresponding index
+        resize: Whether to perform an aspect ratio resize of the patch before the transformations
+        transform: Optional function applied to the image
+        rgb: if False, image will be converted in grayscale
+        channel: 1 or 3. If rgb is True, then channel will be set at 3.
+        balance_classes: if True, the dataset will be balanced by duplicating samples of the minority class
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        samples: list[str],
+        targets: list[str | int],
+        class_to_idx: dict | None = None,
+        resize: int | None = None,
+        transform: Callable | None = None,
+        rgb: bool = True,
+        channel: int = 3,
+        balance_classes: bool = False,
+    ):
+        super().__init__()
+
+        # Keep-Aspect-Ratio resize
+        self.resize = resize
+        self.data_path = data_path
+
+        if balance_classes:
+            samples_array = np.array(samples)
+            targets_array = np.array(targets)
+            samples_to_use: list[str] = []
+            targets_to_use: list[str | int] = []
+
+            cls, counts = np.unique(targets_array, return_counts=True)
+            max_count = np.max(counts)
+            for cl, count in zip(cls, counts):
+                idx_to_pick = list(np.where(targets_array == cl)[0])
+
+                if count < max_count:
+                    idx_to_pick += random.choices(idx_to_pick, k=max_count - count)
+
+                samples_to_use.extend(samples_array[idx_to_pick])
+                targets_to_use.extend(targets_array[idx_to_pick])
+        else:
+            samples_to_use = samples
+            targets_to_use = targets
+
+        # Data
+        self.x = np.array(samples_to_use)
+        self.y = np.array(targets_to_use)
+
+        if class_to_idx is None:
+            unique_targets = np.unique(targets_to_use)
+            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] is not None else None) 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]
+
+        h5_file = h5py.File(path)
+        x = cv2.imread(os.path.join(self.data_path, h5_file["img_path"][()].decode("utf-8")))
+
+        weights = h5_file["triangles_weights"][()]
+        patch_size = h5_file["patch_size"][()]
+
+        if weights.shape[0] == 0:  # pylint: disable=no-member
+            # If the image is completely good sample a point anywhere
+            patch_y = np.random.randint(0, x.shape[0] + 1)
+            patch_x = np.random.randint(0, x.shape[1] + 1)
+        else:
+            random_triangle = np.random.choice(weights.shape[0], p=weights)
+            [patch_y, patch_x] = trisample(h5_file["triangles"][random_triangle])
+
+        h5_file.close()
+
+        # If the patch is outside the image reduce the exceeding area by taking more patch from the inner area
+        [y_left, y_right] = compute_safe_patch_range(patch_y, patch_size[0], x.shape[0])
+        [x_left, x_right] = compute_safe_patch_range(patch_x, patch_size[1], x.shape[1])
+
+        x = x[patch_y - y_left : patch_y + y_right, patch_x - x_left : patch_x + x_right]
+
+        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]
+
+        # 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)

quadra/datasets/segmentation.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from typing import Any
+
+import albumentations
+import cv2
+import numpy as np
+import torch
+
+from quadra.utils.deprecation import deprecated
+from quadra.utils.imaging import keep_aspect_ratio_resize
+from quadra.utils.segmentation import smooth_mask
+
+
+# DEPRECATED -> we can use SegmentationDatasetMulticlass also for one class segmentation
+@deprecated("Use SegmentationDatasetMulticlass instead")
+class SegmentationDataset(torch.utils.data.Dataset):
+    """Custom SegmentationDataset class for loading images and masks.
+
+    Args:
+        image_paths: List of paths to images.
+        mask_paths: List of paths to masks.
+        batch_size: Batch size.
+        object_masks: List of paths to object masks.
+        resize: Resize image to this size.
+        mask_preprocess: Preprocess mask.
+        labels: List of labels.
+        transform: Transformations to apply to images and masks.
+        mask_smoothing: Smooth mask.
+        defect_transform: Transformations to apply to images and masks for defects.
+    """
+
+    def __init__(
+        self,
+        image_paths: list[str],
+        mask_paths: list[str],
+        batch_size: int | None = None,
+        object_masks: list[np.ndarray | Any] | None = None,
+        resize: int = 224,
+        mask_preprocess: Callable | None = None,
+        labels: list[str] | None = None,
+        transform: albumentations.Compose | None = None,
+        mask_smoothing: bool = False,
+        defect_transform: albumentations.Compose | None = None,
+    ):
+        self.transform = transform
+        self.defect_transform = defect_transform
+        self.image_paths = image_paths
+        self.mask_paths = mask_paths
+        self.labels = labels
+        self.mask_preprocess = mask_preprocess
+        self.resize = resize
+        self.object_masks = object_masks
+        self.data_len = len(self.image_paths)
+        self.batch_size = None if batch_size is None else max(batch_size, self.data_len)
+        self.smooth_mask = mask_smoothing
+
+    def __getitem__(self, index):
+        # This is required to avoid infinite loop when running the dataset outside of a dataloader
+        if self.batch_size is not None and self.batch_size == index:
+            raise StopIteration
+
+        if self.batch_size is None and self.data_len == index:
+            raise StopIteration
+
+        index = index % self.data_len
+        image_path = self.image_paths[index]
+
+        image = cv2.imread(str(image_path))
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        object_mask_path = self.object_masks[index] if self.object_masks is not None else None
+        if object_mask_path is not None:
+            object_mask = cv2.imread(str(object_mask_path), 0) if os.path.isfile(object_mask_path) else None
+        else:
+            object_mask = None
+        label = self.labels[index] if self.labels is not None else None
+        if (
+            self.mask_paths[index] is np.nan
+            or self.mask_paths[index] is None
+            or not os.path.isfile(self.mask_paths[index])
+        ):
+            mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
+        else:
+            mask_path = self.mask_paths[index]
+            mask = cv2.imread(str(mask_path), 0)
+        if self.defect_transform is not None and label == 1 and np.sum(mask) == 0:
+            if object_mask is not None:
+                object_mask *= 255
+            aug = self.defect_transform(image=image, mask=mask, object_mask=object_mask, label=label)
+            image = aug["image"]
+            mask = aug["mask"]
+            label = aug["label"]
+        if self.mask_preprocess:
+            mask = self.mask_preprocess(mask)
+            if object_mask is not None:
+                object_mask = self.mask_preprocess(object_mask)
+        if self.resize:
+            image = keep_aspect_ratio_resize(image, self.resize)
+            mask = keep_aspect_ratio_resize(mask, self.resize)
+            if object_mask is not None:
+                object_mask = keep_aspect_ratio_resize(object_mask, self.resize)
+
+        if self.transform is not None:
+            aug = self.transform(image=image, mask=mask)
+            image = aug["image"]
+            mask = aug["mask"]
+        if isinstance(mask, np.ndarray):
+            mask_sum = np.sum(mask)
+        elif isinstance(mask, torch.Tensor):
+            mask_sum = torch.sum(mask)
+        else:
+            raise ValueError("Unsupported type for mask")
+        if mask_sum > 0 and (label is None or label == 0):
+            label = 1
+        if mask_sum == 0:
+            label = 0
+
+        if isinstance(image, np.ndarray):
+            mask = (mask > 0).astype(np.uint8)
+
+            if self.smooth_mask:
+                mask = smooth_mask(mask)
+            mask = np.expand_dims(mask, axis=0)
+        else:
+            mask = (mask > 0).int()
+            if self.smooth_mask:
+                mask = torch.from_numpy(smooth_mask(mask.numpy()))
+            mask = mask.unsqueeze(0)
+
+        return image, mask, label
+
+    def __len__(self):
+        if self.batch_size is None:
+            return self.data_len
+
+        return max(self.data_len, self.batch_size)
+
+
+class SegmentationDatasetMulticlass(torch.utils.data.Dataset):
+    """Custom SegmentationDataset class for loading images and multilabel masks.
+
+    Args:
+        image_paths: List of paths to images.
+        mask_paths: List of paths to masks.
+        idx_to_class: dict with corrispondence btw mask index and classes: {1: class_1, 2: class_2, ..., N: class_N}
+        batch_size: Batch size.
+        transform: Transformations to apply to images and masks.
+        one_hot: if True return a binary mask (n_classxHxW), otherwise the labelled mask HxW. SMP loss requires the
+            second format.
+    """
+
+    def __init__(
+        self,
+        image_paths: list[str],
+        mask_paths: list[str],
+        idx_to_class: dict,
+        batch_size: int | None = None,
+        transform: albumentations.Compose | None = None,
+        one_hot: bool = False,
+    ):
+        self.transform = transform
+        self.image_paths = image_paths
+        self.mask_paths = mask_paths
+        self.idx_to_class = idx_to_class
+        self.data_len = len(self.image_paths)
+        self.batch_size = None if batch_size is None else max(batch_size, self.data_len)
+        self.one_hot = one_hot
+
+    def _preprocess_mask(self, mask: np.ndarray):
+        """Function to preprocess the mask -> needed for albumentations
+        Args:
+            mask: a numpy array of dimension HxW with values in [0] + self.idx_to_class.
+
+        Output:
+            a binary numpy array with dims len(self.idx_to_class) + 1 x H x W
+        """
+        multilayer_mask = np.zeros((len(self.idx_to_class) + 1, *mask.shape[:2]))
+        # provide background information for completeness
+        # single channel mask does not use it anyway.
+        multilayer_mask[0] = (mask == 0).astype(np.uint8)
+        for idx in self.idx_to_class:
+            multilayer_mask[int(idx)] = (mask == int(idx)).astype(np.uint8)
+
+        return multilayer_mask
+
+    def __getitem__(self, index):
+        """Get image and mask."""
+        # This is required to avoid infinite loop when running the dataset outside of a dataloader
+        if self.batch_size is not None and self.batch_size == index:
+            raise StopIteration
+        if self.batch_size is None and self.data_len == index:
+            raise StopIteration
+
+        index = index % self.data_len
+        image_path = self.image_paths[index]
+
+        image = cv2.imread(str(image_path))
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        if (
+            self.mask_paths[index] is np.nan
+            or self.mask_paths[index] is None
+            or not os.path.isfile(self.mask_paths[index])
+        ):
+            mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
+        else:
+            mask_path = self.mask_paths[index]
+            mask = cv2.imread(str(mask_path), 0)
+
+        # we go back to binary masks avoid transformation errors
+        mask = self._preprocess_mask(mask)
+
+        if self.transform is not None:
+            masks = list(mask)
+            aug = self.transform(image=image, masks=masks)
+            image = aug["image"]
+            mask = np.stack(aug["masks"])  # C x H x W
+
+        # we compute single channel mask again
+        # zero is the background
+        if not self.one_hot:  # one hot is done by smp dice loss
+            mask_out = np.zeros(mask.shape[1:])
+            for i in range(1, mask.shape[0]):
+                mask_out[mask[i] == 1] = i
+            # mask_out shape -> HxW
+        else:
+            mask_out = mask
+            # mask_out shape -> CxHxW where C is number of classes (included the background)
+
+        return image, mask_out.astype(int), 0
+
+    def __len__(self):
+        """Returns the dataset lenght."""
+        if self.batch_size is None:
+            return self.data_len
+
+        return max(self.data_len, self.batch_size)

quadra/datasets/ssl.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import random
+from collections.abc import Iterable
+from enum import Enum
+
+import albumentations as A
+import numpy as np
+from torch.utils.data import Dataset
+
+
+class AugmentationStrategy(Enum):
+    """Augmentation Strategy for TwoAugmentationDataset."""
+
+    SAME_IMAGE = 1
+    SAME_CLASS = 2
+
+
+class TwoAugmentationDataset(Dataset):
+    """Two Image Augmentation Dataset for using in self-supervised learning.
+
+    Args:
+        dataset: A torch Dataset object
+        transform: albumentation transformations for each image.
+            If you use single transformation, it will be applied to both images.
+            If you use tuple, it will be applied to first image and second image separately.
+        strategy: Defaults to AugmentationStrategy.SAME_IMAGE.
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        transform: A.Compose | tuple[A.Compose, A.Compose],
+        strategy: AugmentationStrategy = AugmentationStrategy.SAME_IMAGE,
+    ):
+        self.dataset = dataset
+        self.transform = transform
+        self.stategy = strategy
+        if isinstance(transform, Iterable) and not isinstance(transform, str) and len(set(transform)) != 2:
+            raise ValueError("transform must be an Iterable of length 2")
+
+    def __getitem__(self, index):
+        image1, target = self.dataset[index]
+
+        if self.stategy == AugmentationStrategy.SAME_IMAGE:
+            image2 = image1
+        elif self.stategy == AugmentationStrategy.SAME_CLASS:
+            positive_pair_idx = random.choice(np.where(self.dataset.y == target)[0])
+            image2, _ = self.dataset[positive_pair_idx]
+        else:
+            raise ValueError("Unknown strategy")
+
+        if isinstance(self.transform, Iterable):
+            image1 = self.transform[0](image=image1)["image"]
+            image2 = self.transform[1](image=image2)["image"]
+        else:
+            image1 = self.transform(image=image1)["image"]
+            image2 = self.transform(image=image2)["image"]
+
+        return [image1, image2], target
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class TwoSetAugmentationDataset(Dataset):
+    """Two Set Augmentation Dataset for using in self-supervised learning (DINO).
+
+    Args:
+        dataset: Base dataset
+        global_transforms: Global transformations for each image.
+        local_transform: Local transformations for each image.
+        num_local_transforms: Number of local transformations to apply. In total you will have
+            two + num_local_transforms transformations for each image. First element of the array will always
+            return the original image.
+
+    Example:
+        >>> images[0] = global_transform[0](original_image)
+        >>> images[1] = global_transform[1](original_image)
+        >>> images[2:] = local_transform(s)(original_image)
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        global_transforms: tuple[A.Compose, A.Compose],
+        local_transform: A.Compose,
+        num_local_transforms: int,
+    ):
+        self.dataset = dataset
+        self.global_transforms = global_transforms
+        self.local_transform = local_transform
+        self.num_local_transforms = num_local_transforms
+
+        if num_local_transforms < 1:
+            raise ValueError("num_local_transforms must be greater than 0")
+
+    def __getitem__(self, index):
+        original_image, target = self.dataset[index]
+        global_outputs = []
+        local_outputs = []
+        for global_transform in self.global_transforms:
+            global_outputs.append(global_transform(image=original_image)["image"])
+        for _ in range(self.num_local_transforms):
+            local_outputs.append(self.local_transform(image=original_image)["image"])
+        all_outputs = global_outputs + local_outputs
+        return all_outputs, target
+
+    def __len__(self):
+        return len(self.dataset)

quadra/losses/classification/__init__.py

@@ -0,0 +1,6 @@
+from .asl import AsymmetricLoss
+from .focal import FocalLoss
+
+# TODO: Implement prototypical loss as a module
+
+__all__ = ["AsymmetricLoss", "FocalLoss"]

quadra/losses/classification/asl.py

@@ -0,0 +1,83 @@
+import torch
+
+
+class AsymmetricLoss(torch.nn.Module):
+    """Notice - optimized version, minimizes memory allocation and gpu uploading,
+    favors inplace operations.
+
+    Args:
+        gamma_neg: gamma for negative samples
+        gamma_pos: gamma for positive samples
+        m: bias value added to negative samples
+        eps: epsilon to avoid division by zero
+        disable_torch_grad_focal_loss: if True, disables torch grad for focal loss
+        apply_sigmoid: if True, applies sigmoid to input before computing loss
+    """
+
+    def __init__(
+        self,
+        gamma_neg: float = 4,
+        gamma_pos: float = 0,
+        m: float = 0.05,
+        eps: float = 1e-8,
+        disable_torch_grad_focal_loss: bool = False,
+        apply_sigmoid: bool = True,
+    ):
+        super().__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.m = m
+        self.eps = eps
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.apply_sigmoid = apply_sigmoid
+
+        # prevent memory allocation and gpu uploading every iteration, and encourages inplace operations
+        self.targets: torch.Tensor
+        self.anti_targets: torch.Tensor
+        self.xs_pos: torch.Tensor
+        self.xs_neg: torch.Tensor
+        self.asymmetric_w: torch.Tensor
+        self.loss: torch.Tensor
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """Compute the asymmetric loss.
+
+        Args:
+            x: input logits (after sigmoid)
+            y: targets (multi-label binarized vector)
+
+        Returns:
+            asymettric loss
+        """
+        self.targets = y
+        self.anti_targets = 1 - y
+
+        # Calculating Probabilities
+        self.xs_pos = x
+        if self.apply_sigmoid:
+            self.xs_pos = torch.sigmoid(self.xs_pos)
+        self.xs_neg = 1.0 - self.xs_pos
+
+        # Asymmetric clipping
+        if self.m is not None and self.m > 0:
+            self.xs_neg.add_(self.m).clamp_(max=1)
+
+        # Basic CE calculation
+        self.loss = self.targets * torch.log(self.xs_pos.clamp(min=self.eps))
+        self.loss.add_(self.anti_targets * torch.log(self.xs_neg.clamp(min=self.eps)))
+
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(False)
+            self.xs_pos = self.xs_pos * self.targets
+            self.xs_neg = self.xs_neg * self.anti_targets
+            self.asymmetric_w = torch.pow(
+                1 - self.xs_pos - self.xs_neg, self.gamma_pos * self.targets + self.gamma_neg * self.anti_targets
+            )
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(True)
+            self.loss *= self.asymmetric_w
+
+        return -self.loss.sum()