| from cmath import nan |
| import csv |
| import json |
| import logging |
| import os |
| import sys |
| import pydicom |
|
|
| from abc import abstractmethod |
| from itertools import islice |
| from typing import List, Tuple, Dict, Any |
| from torch.utils.data import DataLoader |
| import PIL |
| from torch.utils.data import Dataset |
| import numpy as np |
| import pandas as pd |
| from torchvision import transforms |
| from PIL import Image |
| from skimage import exposure |
| import torch |
|
|
| |
| from torchvision.transforms import InterpolationMode |
|
|
|
|
| class RSNA2018_Dataset(Dataset): |
| def __init__(self, csv_path): |
| data_info = pd.read_csv(csv_path) |
| self.img_path_list = np.asarray(data_info.iloc[:, 1]) |
| self.class_list = np.asarray(data_info.iloc[:, 3]) |
| self.bbox = np.asarray(data_info.iloc[:, 2]) |
| normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) |
|
|
| self.transform = transforms.Compose( |
| [transforms.Resize([224, 224]), transforms.ToTensor(), normalize,] |
| ) |
| self.seg_transfrom = transforms.Compose( |
| [ |
| transforms.ToTensor(), |
| transforms.Resize([224, 224], interpolation=InterpolationMode.NEAREST), |
| ] |
| ) |
|
|
| def __getitem__(self, index): |
| img_path = self.img_path_list[index] |
| class_label = np.array([self.class_list[index]]) |
|
|
| img = self.read_dcm(img_path) |
| image = self.transform(img) |
|
|
| bbox = self.bbox[index] |
| seg_map = np.zeros((1024, 1024)) |
| if class_label == 1: |
| boxes = bbox.split("|") |
| for box in boxes: |
| cc = box.split(";") |
| seg_map[ |
| int(float(cc[1])) : (int(float(cc[1])) + int(float(cc[3]))), |
| int(float(cc[0])) : (int(float(cc[0])) + int(float(cc[2]))), |
| ] = 1 |
| seg_map = self.seg_transfrom(seg_map) |
| return { |
| "image": image, |
| "label": class_label, |
| "image_path": img_path, |
| "seg_map": seg_map, |
| } |
|
|
| def read_dcm(self, dcm_path): |
| dcm_data = pydicom.read_file(dcm_path) |
| img = dcm_data.pixel_array.astype(float) / 255.0 |
| img = exposure.equalize_hist(img) |
|
|
| img = (255 * img).astype(np.uint8) |
| img = PIL.Image.fromarray(img).convert("RGB") |
| return img |
|
|
| def __len__(self): |
| return len(self.img_path_list) |
|
|
|
|
| def create_loader_RSNA( |
| datasets, samplers, batch_size, num_workers, is_trains, collate_fns |
| ): |
| loaders = [] |
| for dataset, sampler, bs, n_worker, is_train, collate_fn in zip( |
| datasets, samplers, batch_size, num_workers, is_trains, collate_fns |
| ): |
| if is_train: |
| shuffle = sampler is None |
| drop_last = True |
| else: |
| shuffle = False |
| drop_last = False |
| loader = DataLoader( |
| dataset, |
| batch_size=bs, |
| num_workers=n_worker, |
| pin_memory=True, |
| sampler=sampler, |
| shuffle=shuffle, |
| collate_fn=collate_fn, |
| drop_last=drop_last, |
| ) |
| loaders.append(loader) |
| return loaders |
|
|
