dataset.py

import os

import numpy as np
import torch
import tqdm
from skimage.io import imread
from torch.utils.data import Dataset

from utils import crop_sample, normalize_volume, pad_sample, resize_sample


class BrainSegmentationDataset(Dataset):
    """Brain MRI dataset for FLAIR abnormality segmentation"""

    in_channels = 3
    out_channels = 1

    def __init__(
        self,
        images_dir,
        transform=None,
        image_size=256,
        subset="train",
        random_sampling=True,
    ):
        assert subset in ["all", "train", "validation"]

        # read images
        volumes = {}
        masks = {}
        self.volume_fnames = {}
        img_cnt = 0
        print(f"reading {subset} images...")
        for dirpath, dirnames, filenames in os.walk(images_dir):
            image_slices = []
            mask_slices = []
            image_names = []
            for filename in sorted(
                filter(lambda f: ".tif" in f, filenames),
                key=lambda x: int(x.split(".")[-2].split("_")[4]),
            ):
                filepath = os.path.join(dirpath, filename)
                if "mask" in filename:
                    mask_slices.append(imread(filepath, as_gray=True))
                else:
                    image_slices.append(imread(filepath))
                    image_names.append(filename)
                    img_cnt += 1
            if image_slices:
                patient_id = dirpath.split("/")[-1]
                volumes[patient_id] = np.array(image_slices[1:-1])
                masks[patient_id] = np.array(mask_slices[1:-1])
                self.volume_fnames[patient_id] = image_names[1:-1]

        self.patients = sorted(volumes)
        self.volume_fnames

        print(f"preprocessing {subset} volumes...")
        # create list of tuples (volume, mask)
        self.volumes = [(volumes[k], masks[k]) for k in self.patients]

        print(f"cropping {subset} volumes...")
        # crop to smallest enclosing volume
        self.volumes = [crop_sample(v) for v in self.volumes]

        print(f"padding {subset} volumes...")
        # pad to square
        self.volumes = [pad_sample(v) for v in self.volumes]

        print(f"resizing {subset} volumes...")
        # resize
        self.volumes = [resize_sample(v, size=image_size) for v in tqdm.tqdm(self.volumes)]

        print(f"normalizing {subset} volumes...")
        # normalize channel-wise
        self.volumes = [(normalize_volume(v), m) for v, m in self.volumes]

        # probabilities for sampling slices based on masks
        self.slice_weights = [m.sum(axis=-1).sum(axis=-1) for v, m in self.volumes]
        self.slice_weights = [
            (s + (s.sum() * 0.1 / len(s))) / (s.sum() * 1.1) for s in self.slice_weights
        ]

        # add channel dimension to masks
        self.volumes = [(v, m[..., np.newaxis]) for (v, m) in self.volumes]

        print(f"done creating {subset} dataset")

        # create global index for patient and slice (idx -> (p_idx, s_idx))
        num_slices = [v.shape[0] for v, m in self.volumes]
        self.patient_slice_index = list(
            zip(
                sum(([i] * num_slices[i] for i in range(len(num_slices))), []),
                sum((list(range(x)) for x in num_slices), []),
            )
        )

        self.random_sampling = random_sampling

        self.transform = transform

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

    def _get_image(self, idx, do_transform=True):
        patient = self.patient_slice_index[idx][0]
        slice_n = self.patient_slice_index[idx][1]

        if self.random_sampling:
            patient = np.random.randint(len(self.volumes))
            slice_n = np.random.choice(
                range(self.volumes[patient][0].shape[0]), p=self.slice_weights[patient]
            )

        v, m = self.volumes[patient]
        image = v[slice_n]
        mask = m[slice_n]
        fname = self.volume_fnames[self.patients[patient]][slice_n]

        if do_transform and self.transform is not None:
            image, mask = self.transform((image, mask))

        # fix dimensions (C, H, W)
        image = image.transpose(2, 0, 1)
        mask = mask.transpose(2, 0, 1)

        image_tensor = torch.from_numpy(image.astype(np.float32)) / 255.0
        mask_tensor = torch.from_numpy(mask.astype(np.float32)) / 255.0

        # return tensors
        return image_tensor, mask_tensor, fname

    def __getitem__(self, idx):
        return self._get_image(idx, do_transform=True)

    def get_original_image(self, idx):
        return self._get_image(idx, do_transform=False)