import os from typing import (Dict, Optional, Tuple, Union, Callable, Iterable) import pandas as pd from PIL import Image from enum import Enum import numpy as np from numpy.random import RandomState import collections.abc from collections import Counter, defaultdict import torch import torch.utils.data as data from torch.utils.data import DataLoader from labelmap import DR_LABELMAP DataRecord = Tuple[Image.Image, int] class RetinopathyDataset(data.Dataset[DataRecord]): """ A class to access the pre-downloaded Diabetic Retinopathy dataset. """ def __init__(self, data_path: str) -> None: """ Constructor. Args: data_path (str): path to the dataset, ex: "retinopathy_data" containing "trainLabels.csv" and "train/". """ super().__init__() self.data_path = data_path self.ext = ".jpeg" anno_path = os.path.join(data_path, "trainLabels.csv") self.anno_df = pd.read_csv(anno_path) # ['image', 'level'] anno_name_set = set(self.anno_df['image']) if True: train_path = os.path.join(data_path, "train") img_path_list = os.listdir(train_path) img_name_set = set([os.path.splitext(p)[0] for p in img_path_list]) assert anno_name_set == img_name_set self.label_map = DR_LABELMAP def __getitem__(self, index: Union[int, slice]) -> DataRecord: assert isinstance(index, int) img_path = self.get_path_at(index) img = Image.open(img_path) label = self.get_label_at(index) return img, label def __len__(self) -> int: return len(self.anno_df) def get_label_at(self, index: int) -> int: label = self.anno_df['level'].iloc[index].item() return label def get_path_at(self, index: int) -> str: img_name = self.anno_df['image'].iloc[index] img_path = os.path.join(self.data_path, "train", img_name+self.ext) return img_path """ Purpose of a split: training or validation. """ class Purpose(Enum): Train = 0 Val = 1 """ Augmentation transformations for an image and a label. """ FeatureAndTargetTransforms = Tuple[Callable[..., torch.Tensor], Callable[..., torch.Tensor]] """ Feature (image) and target (label) tensors. """ TensorRecord = Tuple[torch.Tensor, torch.Tensor] class Split(data.Dataset[TensorRecord], collections.abc.Sequence[TensorRecord]): """ Split is a class that keep a view on a part of a dataset. Split is used to hold the imormation about which samples go to training and which to validation without a need to put these groups of files into separate folders. """ def __init__(self, dataset: RetinopathyDataset, indices: np.ndarray, purpose: Purpose, transforms: FeatureAndTargetTransforms, oversample_factor: int = 1, stratify_classes: bool = False, use_log_frequencies: bool = False, ): """ Constructor. Args: dataset (RetinopathyDataset): The dataset on which the Split "views". indices (np.ndarray): Externally provided indices of samples that are "viewed" on. purpose (Purpose): Either train or val, to be able to replicate the data for train split for effecient workers utilization. transforms (FeatureAndTargetTransforms): Functors of feature and target transforms. oversample_factor (int, optional): Expand the training dataset by replication to avoid dataloader stalls on epoch ends. Defaults to 1. stratify_classes (bool, optional): Whether to apply stratified sampling. Defaults to False. use_log_frequencies (bool, optional): If stratify_classes=True, whether to use logarithmic sampling strategy. If False, apply regular even sampling. Defaults to False. """ self.dataset = dataset self.indices = indices self.purpose = purpose self.feature_transform = transforms[0] self.target_transform = transforms[1] self.oversample_factor = oversample_factor self.stratify_classes = stratify_classes self.use_log_frequencies = use_log_frequencies self.per_class_indices: Optional[Dict[int, np.ndarray]] = None self.frequencies: Optional[Dict[int, float]] = None if self.stratify_classes: self._bucketize_indices() if self.use_log_frequencies: self._calc_frequencies() def _calc_frequencies(self): assert self.per_class_indices is not None counts_dict = {lbl: len(arr) for lbl, arr in self.per_class_indices.items()} counts = np.array(list(counts_dict.values())) counts_nrm = self._normalize(counts) temperature = 50.0 # > 1 to even-out frequencies freqs = self._normalize(np.log1p(counts_nrm * temperature)) self.frequencies = {k: freq.item() for k, freq in zip(self.per_class_indices.keys(), freqs)} print(self.frequencies) @staticmethod def _normalize(arr: np.ndarray) -> np.ndarray: return arr / np.sum(arr) def _bucketize_indices(self): buckets = defaultdict(list) for index in self.indices: label = self.dataset.get_label_at(index) buckets[label].append(index) self.per_class_indices = {k: np.array(v) for k, v in buckets.items()} def __getitem__(self, index: Union[int, slice]) -> TensorRecord: # type: ignore[override] assert isinstance(index, int) if self.purpose == Purpose.Train: index_rem = index % len(self.indices) idx = self.indices[index_rem].item() else: idx = self.indices[index].item() if self.per_class_indices: if self.frequencies is not None: arange = np.arange(len(self.per_class_indices)) frequencies = np.zeros(len(self.per_class_indices), dtype=float) for k, v in self.frequencies.items(): frequencies[k] = v random_key = np.random.choice( arange, p=frequencies) else: random_key = np.random.randint(len(self.per_class_indices)) indices = self.per_class_indices[random_key] actual_index = np.random.choice(indices).item() else: actual_index = idx feature, target = self.dataset[actual_index] feature_tensor = self.feature_transform(feature) target_tensor = self.target_transform(target) return feature_tensor, target_tensor def __len__(self): if self.purpose == Purpose.Train: return len(self.indices) * self.oversample_factor else: return len(self.indices) @staticmethod def make_splits(all_data: RetinopathyDataset, train_transforms: FeatureAndTargetTransforms, val_transforms: FeatureAndTargetTransforms, train_fraction: float, stratify_train: bool, stratify_val: bool, seed: int = 54, ) -> Tuple['Split', 'Split']: """ Prepare train and val splits deterministically. Returns: Tuple[Split, Split]: - Train split - Val split """ prng = RandomState(seed) num_train = int(len(all_data) * train_fraction) all_indices = prng.permutation(len(all_data)) train_indices = all_indices[:num_train] val_indices = all_indices[num_train:] train_data = Split(all_data, train_indices, Purpose.Train, train_transforms, stratify_classes=stratify_train) val_data = Split(all_data, val_indices, Purpose.Val, val_transforms, stratify_classes=stratify_val) return train_data, val_data def print_data_stats(dataset: Union[Iterable[DataRecord], DataLoader], split_name: str) -> None: labels = [] for _, label in dataset: if isinstance(label, torch.Tensor): label = label.cpu().numpy() labels.append(label) labels = np.concatenate(labels) cnt = Counter(labels) print(cnt)