import json
import numpy as np
import torch
from pathlib import Path
import pandas as pd
import pydicom
from ast import literal_eval
from adpkd_segmentation.data.data_utils import (
get_labeled,
get_y_Path,
int16_to_uint8,
make_dcmdicts,
path_2dcm_int16,
path_2label,
TKV_update,
)
from adpkd_segmentation.data.data_utils import (
KIDNEY_PIXELS,
STUDY_TKV,
VOXEL_VOLUME,
)
from adpkd_segmentation.datasets.filters import PatientFiltering
class SegmentationDataset(torch.utils.data.Dataset):
"""Some information about SegmentationDataset"""
def __init__(
self,
label2mask,
dcm2attribs,
patient2dcm,
patient_IDS=None,
augmentation=None,
smp_preprocessing=None,
normalization=None,
output_idx=False,
attrib_types=None,
):
super().__init__()
self.label2mask = label2mask
self.dcm2attribs = dcm2attribs
self.pt2dcm = patient2dcm
self.patient_IDS = patient_IDS
self.augmentation = augmentation
self.smp_preprocessing = smp_preprocessing
self.normalization = normalization
self.output_idx = output_idx
self.attrib_types = attrib_types
# store some attributes as PyTorch tensors
if self.attrib_types is None:
self.attrib_types = {
STUDY_TKV: "float32",
KIDNEY_PIXELS: "float32",
VOXEL_VOLUME: "float32",
}
self.patients = list(patient2dcm.keys())
# kept for compatibility with previous experiments
# following patient order in patient_IDS
if patient_IDS is not None:
self.patients = patient_IDS
self.dcm_paths = []
for p in self.patients:
self.dcm_paths.extend(patient2dcm[p])
self.label_paths = [get_y_Path(dcm) for dcm in self.dcm_paths]
# study_id to TKV and TKV for each dcm
self.studies, self.dcm2attribs = TKV_update(dcm2attribs)
# storring attrib types as tensors
self.tensor_dict = self.prepare_tensor_dict(self.attrib_types)
def __getitem__(self, index):
if isinstance(index, slice):
return [self[ii] for ii in range(*index.indices(len(self)))]
# numpy int16, (H, W)
im_path = self.dcm_paths[index]
image = path_2dcm_int16(im_path)
# image local scaling by default to convert to uint8
if self.normalization is None:
image = int16_to_uint8(image)
else:
image = self.normalization(image, self.dcm2attribs[im_path])
label = path_2label(self.label_paths[index])
# numpy uint8, one hot encoded (C, H, W)
mask = self.label2mask(label[np.newaxis, ...])
if self.augmentation is not None:
# requires (H, W, C) or (H, W)
mask = mask.transpose(1, 2, 0)
sample = self.augmentation(image=image, mask=mask)
image, mask = sample["image"], sample["mask"]
# get back to (C, H, W)
mask = mask.transpose(2, 0, 1)
# convert to float
image = (image / 255).astype(np.float32)
mask = mask.astype(np.float32)
# smp preprocessing requires (H, W, 3)
if self.smp_preprocessing is not None:
image = np.repeat(image[..., np.newaxis], 3, axis=-1)
image = self.smp_preprocessing(image).astype(np.float32)
# get back to (3, H, W)
image = image.transpose(2, 0, 1)
else:
# stack image to (3, H, W)
image = np.repeat(image[np.newaxis, ...], 3, axis=0)
if self.output_idx:
return image, mask, index
return image, mask
def __len__(self):
return len(self.dcm_paths)
def get_verbose(self, index):
"""returns more details than __getitem__()
Args:
index (int): index in dataset
Returns:
tuple: sample, dcm_path, attributes dict
"""
sample = self[index]
dcm_path = self.dcm_paths[index]
attribs = self.dcm2attribs[dcm_path]
return sample, dcm_path, attribs
def get_extra_dict(self, batch_of_idx):
return {k: v[batch_of_idx] for k, v in self.tensor_dict.items()}
def prepare_tensor_dict(self, attrib_types):
tensor_dict = {}
for k, v in attrib_types.items():
tensor_dict[k] = torch.zeros(
self.__len__(), dtype=getattr(torch, v)
)
for idx, _ in enumerate(self):
dcm_path = self.dcm_paths[idx]
attribs = self.dcm2attribs[dcm_path]
for k, v in tensor_dict.items():
v[idx] = attribs[k]
return tensor_dict
class DatasetGetter:
"""Create SegmentationDataset"""
def __init__(
self,
splitter,
splitter_key,
label2mask,
augmentation=None,
smp_preprocessing=None,
filters=None,
normalization=None,
output_idx=False,
attrib_types=None,
):
super().__init__()
self.splitter = splitter
self.splitter_key = splitter_key
self.label2mask = label2mask
self.augmentation = augmentation
self.smp_preprocessing = smp_preprocessing
self.filters = filters
self.normalization = normalization
self.output_idx = output_idx
self.attrib_types = attrib_types
dcms_paths = sorted(get_labeled())
print(
"The number of images before splitting and filtering: {}".format(
len(dcms_paths)
)
)
dcm2attribs, patient2dcm = make_dcmdicts(tuple(dcms_paths))
if filters is not None:
dcm2attribs, patient2dcm = filters(dcm2attribs, patient2dcm)
self.all_patient_IDS = list(patient2dcm.keys())
# train, val, or test
self.patient_IDS = self.splitter(self.all_patient_IDS)[
self.splitter_key
]
patient_filter = PatientFiltering(self.patient_IDS)
self.dcm2attribs, self.patient2dcm = patient_filter(
dcm2attribs, patient2dcm
)
if self.normalization is not None:
self.normalization.update_dcm2attribs(self.dcm2attribs)
def __call__(self):
return SegmentationDataset(
label2mask=self.label2mask,
dcm2attribs=self.dcm2attribs,
patient2dcm=self.patient2dcm,
patient_IDS=self.patient_IDS,
augmentation=self.augmentation,
smp_preprocessing=self.smp_preprocessing,
normalization=self.normalization,
output_idx=self.output_idx,
attrib_types=self.attrib_types,
)
class JsonDatasetGetter:
"""Get the dataset from a prepared patient ID split"""
def __init__(
self,
json_path,
splitter_key,
label2mask,
augmentation=None,
smp_preprocessing=None,
normalization=None,
output_idx=False,
attrib_types=None,
):
super().__init__()
self.label2mask = label2mask
self.augmentation = augmentation
self.smp_preprocessing = smp_preprocessing
self.normalization = normalization
self.output_idx = output_idx
self.attrib_types = attrib_types
dcms_paths = sorted(get_labeled())
print(
"The number of images before splitting and filtering: {}".format(
len(dcms_paths)
)
)
dcm2attribs, patient2dcm = make_dcmdicts(tuple(dcms_paths))
print("Loading ", json_path)
with open(json_path, "r") as f:
dataset_split = json.load(f)
self.patient_IDS = dataset_split[splitter_key]
# filter info dicts to correpsond to patient_IDS
patient_filter = PatientFiltering(self.patient_IDS)
self.dcm2attribs, self.patient2dcm = patient_filter(
dcm2attribs, patient2dcm
)
if self.normalization is not None:
self.normalization.update_dcm2attribs(self.dcm2attribs)
def __call__(self):
return SegmentationDataset(
label2mask=self.label2mask,
dcm2attribs=self.dcm2attribs,
patient2dcm=self.patient2dcm,
patient_IDS=self.patient_IDS,
augmentation=self.augmentation,
smp_preprocessing=self.smp_preprocessing,
normalization=self.normalization,
output_idx=self.output_idx,
attrib_types=self.attrib_types,
)
class InferenceDataset(torch.utils.data.Dataset):
"""Some information about SegmentationDataset"""
def __init__(
self,
dcm2attribs,
patient2dcm,
augmentation=None,
smp_preprocessing=None,
normalization=None,
output_idx=False,
attrib_types=None,
):
super().__init__()
self.dcm2attribs = dcm2attribs
self.pt2dcm = patient2dcm
self.augmentation = augmentation
self.smp_preprocessing = smp_preprocessing
self.normalization = normalization
self.output_idx = output_idx
self.attrib_types = attrib_types
self.patients = list(patient2dcm.keys())
self.dcm_paths = []
for p in self.patients:
self.dcm_paths.extend(patient2dcm[p])
# Sorts Studies by Z axis
studies = [
pydicom.dcmread(path).SeriesDescription for path in self.dcm_paths
]
folders = [path.parent.name for path in self.dcm_paths]
patients = [pydicom.dcmread(path).PatientID for path in self.dcm_paths]
x_dims = [pydicom.dcmread(path).Rows for path in self.dcm_paths]
y_dims = [pydicom.dcmread(path).Columns for path in self.dcm_paths]
z_pos = [
literal_eval(str(pydicom.dcmread(path).ImagePositionPatient))[2]
for path in self.dcm_paths
]
acc_nums = [
pydicom.dcmread(path).AccessionNumber for path in self.dcm_paths
]
ser_nums = [
pydicom.dcmread(path).SeriesNumber for path in self.dcm_paths
]
data = {
"dcm_paths": self.dcm_paths,
"folders": folders,
"studies": studies,
"patients": patients,
"x_dims": x_dims,
"y_dims": y_dims,
"z_pos": z_pos,
"acc_nums": acc_nums,
"ser_nums": ser_nums,
}
group_keys = [
"folders",
"studies",
"patients",
"x_dims",
"y_dims",
"acc_nums",
"ser_nums",
]
dataset = pd.DataFrame.from_dict(data)
dataset["slice_pos"] = ""
grouped_dataset = dataset.groupby(group_keys)
for (name, group) in grouped_dataset:
sort_key = "z_pos"
# handle missing slice position with filename
if group[sort_key].isna().any():
sort_key = "dcm_paths"
zs = list(group[sort_key])
sorted_idxs = np.argsort(zs)
slice_map = {
zs[idx]: pos for idx, pos in zip(sorted_idxs, range(len(zs)))
}
zs_slice_pos = group[sort_key].map(slice_map)
for i in group.index:
dataset.at[i, "slice_pos"] = zs_slice_pos.get(i)
grouped_dataset = dataset.groupby(group_keys)
for (name, group) in grouped_dataset:
group.sort_values(by="slice_pos", inplace=True)
self.df = dataset
self.dcm_paths = list(dataset["dcm_paths"])
def __getitem__(self, index):
if isinstance(index, slice):
return [self[ii] for ii in range(*index.indices(len(self)))]
# numpy int16, (H, W)
im_path = self.dcm_paths[index]
image = path_2dcm_int16(im_path)
# image local scaling by default to convert to uint8
if self.normalization is None:
image = int16_to_uint8(image)
else:
image = self.normalization(image, self.dcm2attribs[im_path])
if self.augmentation is not None:
sample = self.augmentation(image=image)
image = sample["image"]
# convert to float
image = (image / 255).astype(np.float32)
# smp preprocessing requires (H, W, 3)
if self.smp_preprocessing is not None:
image = np.repeat(image[..., np.newaxis], 3, axis=-1)
image = self.smp_preprocessing(image).astype(np.float32)
# get back to (3, H, W)
image = image.transpose(2, 0, 1)
else:
# stack image to (3, H, W)
image = np.repeat(image[np.newaxis, ...], 3, axis=0)
if self.output_idx:
return image, index
return image
def __len__(self):
return len(self.dcm_paths)
def get_verbose(self, index):
"""returns more details than __getitem__()
Args:
index (int): index in dataset
Returns:
tuple: sample, dcm_path, attributes dict
"""
sample = self[index]
dcm_path = self.dcm_paths[index]
attribs = self.dcm2attribs[dcm_path]
return sample, dcm_path, attribs
class InferenceDatasetGetter:
"""Get the dataset from a prepared patient ID split"""
def __init__(
self,
inference_path,
augmentation=None,
smp_preprocessing=None,
normalization=None,
output_idx=False,
attrib_types=None,
):
super().__init__()
self.augmentation = augmentation
self.smp_preprocessing = smp_preprocessing
self.normalization = normalization
self.output_idx = output_idx
self.attrib_types = attrib_types
self.inference_path = Path(inference_path)
all_paths = set(self.inference_path.glob("**/*"))
dcms_paths = []
for path in all_paths:
if path.is_file():
try:
pydicom.filereader.dcmread(path, stop_before_pixels=True)
dcms_paths.append(path)
except pydicom.errors.InvalidDicomError:
continue
self.dcm2attribs, self.patient2dcm = make_dcmdicts(
tuple(dcms_paths), label_status=False, WCM=False
)
if self.normalization is not None:
self.normalization.update_dcm2attribs(self.dcm2attribs)
def __call__(self):
return InferenceDataset(
dcm2attribs=self.dcm2attribs,
patient2dcm=self.patient2dcm,
augmentation=self.augmentation,
smp_preprocessing=self.smp_preprocessing,
normalization=self.normalization,
output_idx=self.output_idx,
attrib_types=self.attrib_types,
)
