-
Notifications
You must be signed in to change notification settings - Fork 23
/
Copy pathdata_loader.py
113 lines (100 loc) · 3.86 KB
/
data_loader.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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import torch.utils.data as data
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import torchvision.transforms as transforms
import os
import numpy as np
from utils import is_img
class Dataset_npy(data.Dataset):
def __init__(self, img_dir):
self.img_list = [os.path.join(img_dir, i) for i in os.listdir(img_dir) if i.endswith(".npy")]
def __getitem__(self, index):
img = np.load(self.img_list[index])
img = Image.fromarray(img).convert("RGB")
img = transforms.RandomCrop(256)(img)
img = transforms.RandomHorizontalFlip()(img)
img = transforms.ToTensor()(img)
return img.squeeze(0), self.img_list[index]
def __len__(self):
return len(self.img_list)
class Dataset(data.Dataset):
def __init__(self, img_dir, shorter_side):
self.img_list = [os.path.join(img_dir, i) for i in os.listdir(img_dir) if is_img(i)]
self.shorter_side = shorter_side
def __getitem__(self, index):
img = Image.open(self.img_list[index]).convert("RGB")
if self.shorter_side:
w, h = img.size
if w < h: # resize the shorter side to `shorter_side`
neww = self.shorter_side
newh = int(h * neww / w)
else:
newh = self.shorter_side
neww = int(w * newh / h)
img = img.resize((neww, newh))
img = transforms.RandomCrop(256)(img)
img = transforms.RandomHorizontalFlip()(img)
img = transforms.ToTensor()(img)
return img.squeeze(0), self.img_list[index]
def __len__(self):
return len(self.img_list)
class TestDataset(data.Dataset):
def __init__(self, img_dir, shorter_side):
self.img_list = [os.path.join(img_dir, i) for i in os.listdir(img_dir) if is_img(i)]
random_order = np.random.permutation(len(self.img_list))
self.img_list = list(np.array(self.img_list)[random_order])
self.shorter_side = shorter_side
def __getitem__(self, index):
img = Image.open(self.img_list[index]).convert("RGB")
if self.shorter_side:
w, h = img.size
if w < h: # resize the shorter side to `shorter_side`
neww = self.shorter_side
newh = int(h * neww / w)
else:
newh = self.shorter_side
neww = int(w * newh / h)
img = img.resize((neww, newh))
img = transforms.CenterCrop(256)(img)
img = transforms.ToTensor()(img)
return img.squeeze(0), self.img_list[index]
def __len__(self):
return len(self.img_list)
class ContentStylePair(data.Dataset):
def __init__(self, pathC, pathS, shorter_side):
self.imgListC = [os.path.join(pathC, i) for i in os.listdir(pathC) if is_img(i)]
self.imgListS = [os.path.join(pathS, i) for i in os.listdir(pathS) if is_img(i)]
self.shorter_side = shorter_side
def __getitem__(self, ix):
imgC = Image.open(self.imgListC[ix % len(self.imgListC)]).convert("RGB")
imgS = Image.open(self.imgListS[ix % len(self.imgListS)]).convert("RGB")
if self.shorter_side:
# content
w, h = imgC.size
if w < h: # resize the shorter side to `shorter_side`
neww = self.shorter_side
newh = int(h * neww / w)
else:
newh = self.shorter_side
neww = int(w * newh / h)
imgC = imgC.resize((neww, newh))
imgC = transforms.RandomCrop(256)(imgC)
imgC = transforms.RandomHorizontalFlip()(imgC)
imgC = transforms.ToTensor()(imgC)
# style
w, h = imgS.size
if w < h: # resize the shorter side to `shorter_side`
neww = self.shorter_side
newh = int(h * neww / w)
else:
newh = self.shorter_side
neww = int(w * newh / h)
imgS = imgS.resize((neww, newh))
imgS = transforms.RandomCrop(256)(imgS)
imgS = transforms.RandomHorizontalFlip()(imgS)
imgS = transforms.ToTensor()(imgS)
return imgC.squeeze(0), imgS.squeeze(0)
def __len__(self):
return max(len(self.imgListC), len(self.imgListS))