File size: 6,552 Bytes
db6ee6a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
import os
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from PIL import Image
from skimage import io
from torch.utils.data import Dataset
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, transforms
from local_config import VIS_ROOT, PATH_TO_MIMIC_CXR
class Chexpert_Dataset(Dataset):
def __init__(self, split='train', truncate=None, loss_weighting="none", use_augs=False):
super().__init__()
# load csv file
self.split = pd.read_csv(f'{PATH_TO_MIMIC_CXR}/mimic-cxr-jpg/2.0.0/mimic-cxr-2.0.0-split.csv')
self.reports = pd.read_csv('mimic-cxr/reports_processed/mimic_cxr_sectioned.csv')
self.reports = self.reports.dropna(subset=['findings'])
self.vis_root = VIS_ROOT
self.img_ids = {img_id: i for i, img_id in enumerate(self.reports['dicom_id'])}
self.split_ids = set(self.split.loc[self.split['split'] == split]['dicom_id'])
self.chexpert = pd.read_csv(f'data/data_files/finding_chexbert_labels.csv')
self.chexpert_cols = ["No Finding", "Enlarged Cardiomediastinum",
"Cardiomegaly", "Lung Opacity",
"Lung Lesion", "Edema",
"Consolidation", "Pneumonia",
"Atelectasis", "Pneumothorax",
"Pleural Effusion", "Pleural Other",
"Fracture", "Support Devices"]
# get all dicom_ids where "split" is split
self.annotation = self.reports.loc[self.reports['dicom_id'].isin(self.split_ids)]
self.annotation['study_id'] = self.annotation['Note_file'].apply(lambda x: int(x.lstrip('s').rstrip('.txt')))
# merge chexpert labels
self.annotation = pd.merge(self.annotation, self.chexpert, how='left', left_on=['dicom_id'], right_on=['dicom_id'])
if truncate is not None:
self.annotation = self.annotation[:truncate]
self.vis_transforms = Compose([Resize(512), CenterCrop(488), ToTensor(), ExpandChannels()])
if use_augs:
aug_tfm = transforms.Compose([transforms.RandomAffine(degrees=30, shear=15),
transforms.ColorJitter(brightness=0.2, contrast=0.2)])
self.vis_transforms = transforms.Compose([self.vis_transforms, aug_tfm])
self.loss_weighting = loss_weighting
def get_class_weights(self):
"""Compute class weights based on the inverse of class frequencies.
Returns:
Dict[str, float]: Class weights.
"""
if self.loss_weighting == "none":
return torch.ones(len(self.chexpert_cols), dtype=torch.float32)
label_counts = torch.zeros(len(self.chexpert_cols), dtype=torch.float32)
# iterate over dataframe getting rows
for _, ann in self.annotation.iterrows():
chexpert_labels = self._extract_chexpert_labels_from_row(ann)
label_counts += chexpert_labels
# Compute class weights
if self.loss_weighting == "lin":
class_weights = len(self.annotation) / label_counts
elif self.loss_weighting == "log":
class_weights = torch.log(len(self.annotation) / label_counts)
return class_weights
def remap_to_uint8(self, array: np.ndarray, percentiles=None) -> np.ndarray:
"""Remap values in input so the output range is :math:`[0, 255]`.
Percentiles can be used to specify the range of values to remap.
This is useful to discard outliers in the input data.
:param array: Input array.
:param percentiles: Percentiles of the input values that will be mapped to ``0`` and ``255``.
Passing ``None`` is equivalent to using percentiles ``(0, 100)`` (but faster).
:returns: Array with ``0`` and ``255`` as minimum and maximum values.
"""
array = array.astype(float)
if percentiles is not None:
len_percentiles = len(percentiles)
if len_percentiles != 2:
message = (
'The value for percentiles should be a sequence of length 2,'
f' but has length {len_percentiles}'
)
raise ValueError(message)
a, b = percentiles
if a >= b:
raise ValueError(f'Percentiles must be in ascending order, but a sequence "{percentiles}" was passed')
if a < 0 or b > 100:
raise ValueError(f'Percentiles must be in the range [0, 100], but a sequence "{percentiles}" was passed')
cutoff: np.ndarray = np.percentile(array, percentiles)
array = np.clip(array, *cutoff)
array -= array.min()
array /= array.max()
array *= 255
return array.astype(np.uint8)
def load_image(self, path) -> Image.Image:
"""Load an image from disk.
The image values are remapped to :math:`[0, 255]` and cast to 8-bit unsigned integers.
:param path: Path to image.
:returns: Image as ``Pillow`` ``Image``.
"""
# Although ITK supports JPEG and PNG, we use Pillow for consistency with older trained models
if path.suffix in [".jpg", ".jpeg", ".png"]:
image = io.imread(path)
else:
raise ValueError(f"Image type not supported, filename was: {path}")
image = self.remap_to_uint8(image)
return Image.fromarray(image).convert("L")
def _extract_chexpert_labels_from_row(self, row: pd.Series) -> torch.Tensor:
labels = torch.zeros(len(self.chexpert_cols), dtype=torch.float32)
for i, col in enumerate(self.chexpert_cols):
if row[col] == 1:
labels[i] = 1
return labels
def __getitem__(self, index):
ann = self.annotation.iloc[index]
image_path = os.path.join(self.vis_root, ann["Img_Folder"], ann["Img_Filename"])
image = self.load_image(Path(image_path))
image = self.vis_transforms(image)
chexpert_labels = self._extract_chexpert_labels_from_row(ann)
return {
"image": image,
"labels": chexpert_labels,
"image_id": self.img_ids[ann["dicom_id"]],
"report": ann["findings"],
"study_id": ann["study_id"],
"dicom_id": ann["dicom_id"],
}
def __len__(self):
return len(self.annotation)
if __name__ == '__main__':
dataset = Chexpert_Dataset()
print(dataset[0])
|