src/datamodules/focus_datamodule.py

import os
from typing import List, Optional, Tuple
import pandas as pd
from skimage import io

import numpy as np

import torch
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, Dataset, random_split
from torchvision.transforms import transforms


class FocusDataSet(Dataset):
    """Dataset for z-stacked images of neglected tropical diseaeses."""

    def __init__(
        self, csv_file, root_dir, transform=None, in_memory=True, additional_col_list=[]
    ):
        """Initialize focus satck dataset.

        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.metadata = pd.read_csv(csv_file)
        self.in_memory = in_memory

        self.additional_col_index = {}

        _col_list = list(additional_col_list)  # clone list to avoid modifying default
        for attribute in _col_list:
            self.additional_col_index[attribute] = self.metadata.columns.get_loc(
                attribute
            )

        self.col_index_path = self.metadata.columns.get_loc("image_path")
        self.col_index_focus = self.metadata.columns.get_loc("focus_height")
        self.root_dir = root_dir
        self.transform = transform

        self.images = []
        if self.in_memory:
            self.images = np.array(
                list(map(self._load_img, self.metadata["image_path"].tolist()))
            )

    def _load_img(self, img_path):
        path = os.path.join(self.root_dir, img_path)
        img = io.imread(path)
        return img

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

        Returns:
            int: the length
        """
        return len(self.metadata)

    def __getitem__(self, idx):
        """Get one items from the dataset.

        Args:
            idx (int) The index of the sample that is to be retrieved

        Returns:
            Item/Items which is a dictionary containing "image" and "focus_height"
        """
        if torch.is_tensor(idx):
            idx = idx.tolist()

        if self.in_memory:
            image = self.images[idx]
        else:
            image = self._load_img(self.metadata.iloc[idx, self.col_index_path])

        if self.transform:
            image = self.transform(image)

        focus_height = torch.from_numpy(
            np.asarray(self.metadata.iloc[idx, self.col_index_focus])
        ).float()

        sample = {"image": image, "focus_height": focus_height}

        for attr, col_idx in self.additional_col_index.items():
            sample[attr] = self.metadata.iloc[idx, col_idx]

        return sample


class FocusDataModule(LightningDataModule):
    """
    LightningDataModule for FocusStack dataset.
    """

    def __init__(
        self,
        data_dir: str = "data/",
        csv_train_file: str = "data/train_metadata.csv",
        csv_val_file: str = "data/validation_metadata.csv",
        csv_test_file: str = "data/test_metadata.csv",
        batch_size: int = 64,
        num_workers: int = 0,
        pin_memory: bool = False,
        in_memory: bool = True,
        augmentation: bool = False,
        additional_col_list: List[str] = [],
    ):
        super().__init__()

        # this line allows to access init params with 'self.hparams' attribute
        self.save_hyperparameters(logger=False)

        transform_list = [
            transforms.ToTensor(),
            transforms.ConvertImageDtype(torch.float),
        ]

        self.base_transforms = []
        self.base_transforms.extend(transform_list)
        self.base_transforms = transforms.Compose(self.base_transforms)

        if augmentation:
            transform_list.extend(
                [
                    transforms.RandomHorizontalFlip(p=0.5),
                    transforms.RandomVerticalFlip(p=0.5),
                    transforms.RandomChoice(
                        [
                            transforms.RandomApply(
                                [transforms.RandomRotation((90, 90))], p=0.5
                            ),
                            transforms.RandomApply(
                                [transforms.RandomRotation((180, 180))], p=0.5
                            ),
                            transforms.RandomApply(
                                [transforms.RandomRotation((270, 270))], p=0.5
                            ),
                        ]
                    ),
                ]
            )

        # data transformations
        self.transforms = transforms.Compose(transform_list)

        self.data_train: Optional[Dataset] = None
        self.data_val: Optional[Dataset] = None
        self.data_test: Optional[Dataset] = None
        self.in_memory = in_memory
        self.additional_col_list = additional_col_list

    def prepare_data(self):
        """This method is not implemented as of yet.

        Download data if needed. This method is called only from a single GPU.
        Do not use it to assign state (self.x = y).
        """
        pass

    def setup(self, stage: Optional[str] = None):
        """Load data. Set variables: `self.data_train`, `self.data_val`, `self.data_test`.
        This method is called by lightning twice for `trainer.fit()` and `trainer.test()`, so be careful if you do a random split!
        The `stage` can be used to differentiate whether it's called before trainer.fit()` or `trainer.test()`."""

        # load datasets only if they're not loaded already
        if not self.data_train and not self.data_val and not self.data_test:
            self.data_train = FocusDataSet(
                self.hparams.csv_train_file,
                self.hparams.data_dir,
                transform=self.transforms,
                in_memory=self.in_memory,
                additional_col_list=self.additional_col_list,
            )

            self.data_val = FocusDataSet(
                self.hparams.csv_val_file,
                self.hparams.data_dir,
                transform=self.base_transforms,
                in_memory=self.in_memory,
                additional_col_list=self.additional_col_list,
            )

            self.data_test = FocusDataSet(
                self.hparams.csv_test_file,
                self.hparams.data_dir,
                transform=self.base_transforms,
                in_memory=self.in_memory,
                additional_col_list=self.additional_col_list,
            )

    def train_dataloader(self):
        return DataLoader(
            dataset=self.data_train,
            batch_size=self.hparams.batch_size,
            num_workers=self.hparams.num_workers,
            pin_memory=self.hparams.pin_memory,
            shuffle=True,
        )

    def val_dataloader(self):
        return DataLoader(
            dataset=self.data_val,
            batch_size=self.hparams.batch_size,
            num_workers=self.hparams.num_workers,
            pin_memory=self.hparams.pin_memory,
            shuffle=False,
        )

    def test_dataloader(self):
        return DataLoader(
            dataset=self.data_test,
            batch_size=self.hparams.batch_size,
            num_workers=self.hparams.num_workers,
            pin_memory=self.hparams.pin_memory,
            shuffle=False,
        )