/
dataloder.py
51 lines (48 loc) · 1.79 KB
/
dataloder.py
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
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from data_loaders.makeup_dataloader import MakeupDataloader
import torch
import numpy as np
import PIL
from psgan.preprocess import PreProcess
def ToTensor(pic):
# handle PIL Image
if pic.mode == 'I':
img = torch.from_numpy(np.array(pic, np.int32, copy=False))
elif pic.mode == 'I;16':
img = torch.from_numpy(np.array(pic, np.int16, copy=False))
else:
img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
# PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
if pic.mode == 'YCbCr':
nchannel = 3
elif pic.mode == 'I;16':
nchannel = 1
else:
nchannel = len(pic.mode)
img = img.view(pic.size[1], pic.size[0], nchannel)
# put it from HWC to CHW format
# yikes, this transpose takes 80% of the loading time/CPU
img = img.transpose(0, 1).transpose(0, 2).contiguous()
if isinstance(img, torch.ByteTensor):
return img.float()
else:
return img
def get_loader(config, mode="train"):
# return the DataLoader
transform = transforms.Compose([
transforms.Resize(config.DATA.IMG_SIZE),
transforms.ToTensor(),
transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])])
transform_mask = transforms.Compose([
transforms.Resize(config.DATA.IMG_SIZE, interpolation=PIL.Image.NEAREST),
ToTensor])
dataset = MakeupDataloader(
config.DATA.PATH, transform=transform,
transform_mask=transform_mask, preprocess=PreProcess(config))
#"""
dataloader = DataLoader(dataset=dataset,
batch_size=config.DATA.BATCH_SIZE,
shuffle=True, num_workers=config.DATA.NUM_WORKERS)
return dataloader