| | 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 |
| |
|
