datasets

AnomalyDataset(transform, samples, task='segmentation', valid_area_mask=None, crop_area=None)

Bases: Dataset

Anomaly Dataset.

Parameters:

• transform (alb.Compose) –

  Albumentations compose.

• task (str) –

  classification or segmentation

• samples (DataFrame) –

  Pandas dataframe containing samples following the same structure created by make_anomaly_dataset

• valid_area_mask (Optional[str]) –

  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[int, int, int, int]]) –

  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.

Source code in quadra/datasets/anomaly.py

def __init__(
    self,
    transform: alb.Compose,
    samples: DataFrame,
    task: str = "segmentation",
    valid_area_mask: Optional[str] = None,
    crop_area: Optional[Tuple[int, int, int, int]] = 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: Optional[np.ndarray] = 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

__getitem__(index)

Get dataset item for the index index.

Parameters:

• index (int) –

  Index to get the item.

Returns:

• Dict[str, Union[str, Tensor]] –

  Dict of image tensor during training.

• Dict[str, Union[str, Tensor]] –

  Otherwise, Dict containing image path, target path, image tensor, label and transformed bounding box.

Source code in quadra/datasets/anomaly.py

def __getitem__(self, index: int) -> Dict[str, Union[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, Union[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])
            else:
                if os.path.isfile(mask_path):
                    mask = cv2.imread(mask_path, flags=0) / 255.0
                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

__len__()

Get length of the dataset.

Source code in quadra/datasets/anomaly.py

def __len__(self) -> int:
    """Get length of the dataset."""
    return len(self.samples)

ClassificationDataset(samples, targets, class_to_idx=None, resize=None, roi=None, transform=None, rgb=True, channel=3, random_padding=False, circular_crop=False)

Bases: ImageClassificationListDataset

Custom Classification Dataset.

Parameters:

• samples (List[str]) –

  List of paths to images

• targets (List[Union[str, int]]) –

  List of targets

• class_to_idx (Optional[Dict]) –

  Defaults to None.

• resize (Optional[int]) –

  Resize image to this size. Defaults to None.

• roi (Optional[Tuple[int, int, int, int]]) –

  Region of interest. Defaults to None.

• transform (Optional[Callable]) –

  transform function. Defaults to None.

• rgb (bool) –

  Use RGB space

• channel (int) –

  Number of channels. Defaults to 3.

• random_padding (bool) –

  Random padding. Defaults to False.

• circular_crop (bool) –

  Circular crop. Defaults to False.

Source code in quadra/datasets/classification.py

def __init__(
    self,
    samples: List[str],
    targets: List[Union[str, int]],
    class_to_idx: Optional[Dict] = None,
    resize: Optional[int] = None,
    roi: Optional[Tuple[int, int, int, int]] = None,
    transform: Optional[Callable] = 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

ImageClassificationListDataset(samples, targets, class_to_idx=None, resize=None, roi=None, transform=None, rgb=True, channel=3, allow_missing_label=False)

Bases: Dataset

Standard classification dataset.

Parameters:

• samples (List[str]) –

  List of paths to images to be read

• targets (List[Union[str, int]]) –

  List of labels, one for every image in samples

• class_to_idx (Optional[Dict]) –

  mapping from classes to unique indexes. Defaults to None.

• resize (Optional[int]) –

  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[Tuple[int, int, int, int]]) –

  Optional ROI, with (x_upper_left, y_upper_left, x_bottom_right, y_bottom_right). Defaults to None.

• transform (Optional[Callable]) –

  Optional Albumentations transform. Defaults to None.

• rgb (bool) –

  if False, image will be converted in grayscale

• channel (int) –

  1 or 3. If rgb is True, then channel will be set at 3.

• allow_missing_label (Optional[bool]) –

  If set to false warn the user if the dataset contains missing labels

Source code in quadra/datasets/classification.py

def __init__(
    self,
    samples: List[str],
    targets: List[Union[str, int]],
    class_to_idx: Optional[Dict] = None,
    resize: Optional[int] = None,
    roi: Optional[Tuple[int, int, int, int]] = None,
    transform: Optional[Callable] = None,
    rgb: bool = True,
    channel: int = 3,
    allow_missing_label: Optional[bool] = 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,
        )

    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

MultilabelClassificationDataset(samples, targets, class_to_idx=None, transform=None, rgb=True)

Bases: torch.utils.data.Dataset

Custom MultilabelClassification Dataset.

Parameters:

• samples (List[str]) –

  list of paths to images.

• targets (np.ndarray) –

  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 (Optional[Dict]) –

  Defaults to None.

• transform (Optional[Callable]) –

  transform function. Defaults to None.

• rgb (bool) –

  Use RGB space

Source code in quadra/datasets/classification.py

def __init__(
    self,
    samples: List[str],
    targets: np.ndarray,
    class_to_idx: Optional[Dict] = None,
    transform: Optional[Callable] = 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

PatchSklearnClassificationTrainDataset(data_path, samples, targets, class_to_idx=None, resize=None, transform=None, rgb=True, channel=3, balance_classes=False)

Bases: 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.

Parameters:

• data_path (str) –

  base path to the dataset

• samples (List[str]) –

  Paths to h5 files

• targets (List[Union[str, int]]) –

  Labels associated with each sample

• class_to_idx (Optional[Dict]) –

  Mapping between class and corresponding index

• resize (Optional[int]) –

  Whether to perform an aspect ratio resize of the patch before the transformations

• transform (Optional[Callable]) –

  Optional function applied to the image

• rgb (bool) –

  if False, image will be converted in grayscale

• channel (int) –

  1 or 3. If rgb is True, then channel will be set at 3.

• balance_classes (bool) –

  if True, the dataset will be balanced by duplicating samples of the minority class

Source code in quadra/datasets/patch.py

def __init__(
    self,
    data_path: str,
    samples: List[str],
    targets: List[Union[str, int]],
    class_to_idx: Optional[Dict] = None,
    resize: Optional[int] = None,
    transform: Optional[Callable] = 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[Union[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

SegmentationDataset(image_paths, mask_paths, batch_size=None, object_masks=None, resize=224, mask_preprocess=None, labels=None, transform=None, mask_smoothing=False, defect_transform=None)

Bases: torch.utils.data.Dataset

Custom SegmentationDataset class for loading images and masks.

Parameters:

• image_paths (List[str]) –

  List of paths to images.

• mask_paths (List[str]) –

  List of paths to masks.

• batch_size (Optional[int]) –

  Batch size.

• object_masks (Optional[List[Union[np.ndarray, Any]]]) –

  List of paths to object masks.

• resize (int) –

  Resize image to this size.

• mask_preprocess (Optional[Callable]) –

  Preprocess mask.

• labels (Optional[List[str]]) –

  List of labels.

• transform (Optional[albumentations.Compose]) –

  Transformations to apply to images and masks.

• mask_smoothing (bool) –

  Smooth mask.

• defect_transform (Optional[albumentations.Compose]) –

  Transformations to apply to images and masks for defects.

Source code in quadra/datasets/segmentation.py

def __init__(
    self,
    image_paths: List[str],
    mask_paths: List[str],
    batch_size: Optional[int] = None,
    object_masks: Optional[List[Union[np.ndarray, Any]]] = None,
    resize: int = 224,
    mask_preprocess: Optional[Callable] = None,
    labels: Optional[List[str]] = None,
    transform: Optional[albumentations.Compose] = None,
    mask_smoothing: bool = False,
    defect_transform: Optional[albumentations.Compose] = 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

SegmentationDatasetMulticlass(image_paths, mask_paths, idx_to_class, batch_size=None, transform=None, one_hot=False)

Bases: torch.utils.data.Dataset

Custom SegmentationDataset class for loading images and multilabel masks.

Parameters:

• image_paths (List[str]) –

  List of paths to images.

• mask_paths (List[str]) –

  List of paths to masks.

• idx_to_class (Dict) –

  dict with corrispondence btw mask index and classes: {1: class_1, 2: class_2, ..., N: class_N}

• batch_size (Optional[int]) –

  Batch size.

• transform (Optional[albumentations.Compose]) –

  Transformations to apply to images and masks.

• one_hot (bool) –

  if True return a binary mask (n_classxHxW), otherwise the labelled mask HxW. SMP loss requires the second format.

Source code in quadra/datasets/segmentation.py

def __init__(
    self,
    image_paths: List[str],
    mask_paths: List[str],
    idx_to_class: Dict,
    batch_size: Optional[int] = None,
    transform: Optional[albumentations.Compose] = 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

__getitem__(index)

Get image and mask.

Source code in quadra/datasets/segmentation.py

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

__len__()

Returns the dataset lenght.

Source code in quadra/datasets/segmentation.py

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)

TwoAugmentationDataset(dataset, transform, strategy=AugmentationStrategy.SAME_IMAGE)

Bases: Dataset

Two Image Augmentation Dataset for using in self-supervised learning.

Parameters:

• dataset (Dataset) –

  A torch Dataset object

• transform (Union[A.Compose, Tuple[A.Compose, A.Compose]]) –

  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 (AugmentationStrategy) –

  Defaults to AugmentationStrategy.SAME_IMAGE.

Source code in quadra/datasets/ssl.py

def __init__(
    self,
    dataset: Dataset,
    transform: Union[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")

TwoSetAugmentationDataset(dataset, global_transforms, local_transform, num_local_transforms)

Bases: Dataset

Two Set Augmentation Dataset for using in self-supervised learning (DINO).

Parameters:

• dataset (Dataset) –

  Base dataset

• global_transforms (Tuple[A.Compose, A.Compose]) –

  Global transformations for each image.

• local_transform (A.Compose) –

  Local transformations for each image.

• num_local_transforms (int) –

  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. images[0] = global_transform0 images[1] = global_transform1 images[2:] = local_transform(s)(original_image) ...

Source code in quadra/datasets/ssl.py

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