dataset.py

import os
from pathlib import Path, PurePath

import numpy as np
import pandas as pd
import SimpleITK as sitk
import torch
from torch.utils.data import Dataset


from dicom_utils.dicom_utils import augmentation as aug
from dicom_utils.dicom_utils import processing as dup


def augment_3D(image, mode, size):
    N_CHANNELS = image.shape[0]
    image_seq = [sitk.GetImageFromArray(image[i, :, :, :]) for i in range(N_CHANNELS)]

    if mode == 'train':
        # morphological augmentation
        image_seq = aug.augment_morph(image_seq)

        # add gaussian noise
        fg = aug.get_gauss_noise()
        image_seq = [fg.Execute(image) for image in image_seq]

    image_seq = [dup.resample(image, size=(size, size, size)) for image in image_seq]
    image_seq = [sitk.GetArrayFromImage(image) for image in image_seq]
    image = np.stack(image_seq, axis=0)
    return image


def augment_3D_HN(image, mode, size):
    def normalize_range(image, range_pixel):
        image = sitk.Threshold(
            image, lower=-5000, upper=range_pixel[1], outsideValue=range_pixel[1]
        )
        image = sitk.Threshold(
            image, lower=range_pixel[0], upper=5000, outsideValue=range_pixel[0]
        )
        image = (image - range_pixel[0]) / (range_pixel[1] - range_pixel[0])
        return image

    image_CT = sitk.GetImageFromArray(image[0, :, :, :])
    image_PT = sitk.GetImageFromArray(image[1, :, :, :])

    # normalize_range
    image_CT = normalize_range(image_CT, [-1000, 3000])
    image_PT = normalize_range(image_PT, [0, 50])

    image_seq = [image_CT, image_PT]

    if mode == 'train':
        # morphological augmentation
        image_seq = aug.augment_morph(image_seq)

        # add gaussian noise
        # fg = aug.get_gauss_noise()
        # image_seq = [fg.Execute(image) for image in image_seq]

    image_seq = [dup.resample(image, size=(size, size, size)) for image in image_seq]
    image_seq = [sitk.GetArrayFromImage(image) for image in image_seq]
    image = np.stack(image_seq, axis=0)
    return image


class NumpyCSVDataset(Dataset):
    def __init__(
        self,
        data_dir,
        clinical_file,
        label_col,
        size,
        transforms=augment_3D,
        mode='train',
        seed=1234,
    ):
        super(NumpyCSVDataset, self).__init__()

        if not isinstance(data_dir, PurePath):
            data_dir = Path(data_dir)

        self.data_dir = data_dir
        self.size = size
        self.transforms = transforms
        self.mode = mode
        self.seed = seed

        np.random.seed(self.seed)

        clinical = pd.read_csv(clinical_file).sort_values(by=['patient'])

        available_files = [f for f in os.listdir(data_dir) if f.endswith('.npy')]

        # filter the clinical file in order to keep files that are really on disk
        clinical = clinical.loc[clinical['filename'].isin(available_files)]

        self._filenames = clinical['filename'].values
        self._patients = clinical['patient'].values
        self._labels = clinical[label_col].values

        self.indices = np.arange(len(self._filenames))

    def __getitem__(self, idx):
        label = self._labels[self.indices[idx]]
        filename = self._filenames[self.indices[idx]]
        patient = self._patients[self.indices[idx]]
        data_file = self.data_dir / self._filenames[self.indices[idx]]
        data = np.load(data_file)

        data = self.transforms(data, self.mode, self.size)
        data = torch.Tensor(data)
        output = {
            'data': data,
            'target': label,
            'filename': filename,
            'patient': patient,
        }
        return output

    @property
    def labels(self):
        return self._labels[self.indices]

    @property
    def filenames(self):
        return self._filenames[self.indices]

    @property
    def patients(self):
        return self._patients[self.indices]

    def __len__(self):
        return len(self.indices)

    def __shuffle__(self):
        idx_permut = np.random.permutation(self.__len__())
        self._filenames = self._filenames[idx_permut]
        self._labels = self._labels[idx_permut]
        self._patients = self._patients[idx_permut]
        self.indices = self.indices[idx_permut]