Spaces:
Running
Running
from PIL import Image | |
from torch.utils.data import Dataset, DataLoader | |
from augmentations import ha_augment_sample, resize_sample, spatial_augment_sample | |
from utils import to_tensor_sample | |
def image_transforms(shape, jittering): | |
def train_transforms(sample): | |
sample = resize_sample(sample, image_shape=shape) | |
sample = spatial_augment_sample(sample) | |
sample = to_tensor_sample(sample) | |
sample = ha_augment_sample(sample, jitter_paramters=jittering) | |
return sample | |
return {'train': train_transforms} | |
class GetData(Dataset): | |
def __init__(self, config, transforms=None): | |
""" | |
Get the list containing all images and labels. | |
""" | |
datafile = open(config.train_txt, 'r') | |
lines = datafile.readlines() | |
dataset = [] | |
for line in lines: | |
line = line.rstrip() | |
data = line.split() | |
dataset.append(data[0]) | |
self.config = config | |
self.dataset = dataset | |
self.root = config.train_root | |
self.transforms = transforms | |
def __getitem__(self, index): | |
""" | |
Return image'data and its label. | |
""" | |
img_path = self.dataset[index] | |
img_file = self.root + img_path | |
img = Image.open(img_file) | |
# image.mode == 'L' means the image is in gray scale | |
if img.mode == 'L': | |
img_new = Image.new("RGB", img.size) | |
img_new.paste(img) | |
sample = {'image': img_new, 'idx': index} | |
else: | |
sample = {'image': img, 'idx': index} | |
if self.transforms: | |
sample = self.transforms(sample) | |
return sample | |
def __len__(self): | |
""" | |
Return the number of all data. | |
""" | |
return len(self.dataset) | |
def get_data_loader( | |
config, | |
transforms=None, | |
sampler=None, | |
drop_last=True, | |
): | |
""" | |
Return batch data for training. | |
""" | |
transforms = image_transforms(shape=config.image_shape, jittering=config.jittering) | |
dataset = GetData(config, transforms=transforms['train']) | |
train_loader = DataLoader( | |
dataset, | |
batch_size=config.batch_size, | |
shuffle=config.shuffle, | |
sampler=sampler, | |
num_workers=config.num_workers, | |
pin_memory=config.pin_memory, | |
drop_last=drop_last | |
) | |
return train_loader | |