| import csv |
| import json |
| import logging |
| import os |
| import sys |
| 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 dataset.randaugment import RandomAugment |
|
|
|
|
| class Chestxray14_Dataset(Dataset): |
| def __init__(self, csv_path, is_train=True): |
| data_info = pd.read_csv(csv_path) |
|
|
| self.img_path_list = np.asarray(data_info.iloc[:, 0]) |
| self.class_list = np.asarray(data_info.iloc[:, 2:]) |
|
|
| normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) |
| if is_train: |
| self.transform = transforms.Compose( |
| [ |
| transforms.RandomResizedCrop( |
| 224, scale=(0.2, 1.0), interpolation=Image.BICUBIC |
| ), |
| transforms.RandomHorizontalFlip(), |
| RandomAugment( |
| 2, |
| 7, |
| isPIL=True, |
| augs=[ |
| "Identity", |
| "AutoContrast", |
| "Equalize", |
| "Brightness", |
| "Sharpness", |
| "ShearX", |
| "ShearY", |
| "TranslateX", |
| "TranslateY", |
| "Rotate", |
| ], |
| ), |
| transforms.ToTensor(), |
| normalize, |
| ] |
| ) |
| else: |
| self.transform = transforms.Compose( |
| [transforms.Resize([224, 224]), transforms.ToTensor(), normalize,] |
| ) |
|
|
| def __getitem__(self, index): |
| img_path = self.img_path_list[index] |
| class_label = self.class_list[index] |
| img = PIL.Image.open(img_path).convert("RGB") |
| image = self.transform(img) |
|
|
| return {"image": image, "label": class_label} |
|
|
| def __len__(self): |
| return len(self.img_path_list) |
|
|
|
|
| def create_loader(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 |
|
|
