-
Notifications
You must be signed in to change notification settings - Fork 0
/
my_dataset.py
118 lines (77 loc) · 3.26 KB
/
my_dataset.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
114
115
116
117
118
import torch
import random
import os
from torch.utils.data import DataLoader
import torchvision
from torchvision import transforms
from PIL import Image
class MyDataset(torch.utils.data.Dataset):
"""Custom dataset."""
def __init__(self, root_dir, transform=None):
"""
Args:
root_dir (string): Directory with train/val/test the images.
transform (callable, optional): Optional transform to be applied on a sample.
"""
self.root_dir = root_dir
self.transform = transform
def __len__(self):
return len(os.listdir(path=self.root_dir))
def __getitem__(self, idx):
img_name = os.path.join(self.root_dir, str(idx+1)+'.jpg')
image = Image.open(img_name)
image = torchvision.transforms.functional.to_tensor(image)
width = int(image.size(-1) / 2)
image, target = image[:, :, width:], image[:, :, :width]
if self.transform:
image, target = self.transform(image, target)
sample = {'x': image, 'target': target}
return sample
def my_transforms_tr(sample, target):
# normalize for right pixel values; for inference use backward op: *0.5 + 0.5
sample = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(sample)
target = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(target)
x = random.randint(0, 30)
y = random.randint(0, 30)
flip = random.random() > 0.5
sample = transforms.Resize([286, 286], Image.BICUBIC)(sample)
target = transforms.Resize([286, 286], Image.BICUBIC)(target)
#sample = transforms.RandomCrop(256)(sample)
sample = sample[:, x:256 + x, y:256 + y]
target = target[:, x:256 + x, y:256 + y]
if flip:
actual_flip = transforms.RandomHorizontalFlip()
sample = actual_flip(sample)
target = actual_flip(target)
return sample, target
def my_transforms_val(sample, target):
sample = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(sample)
target = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(target)
return sample, target
### for day2night dataset
class MyDataset_day2night(torch.utils.data.Dataset):
"""Custom dataset."""
def __init__(self, root_dir, transform=None):
"""
Args:
root_dir (string): Directory with train/val/test the images.
transform (callable, optional): Optional transform to be applied on a sample.
"""
self.root_dir = root_dir
self.transform = transform
self.df = pd.DataFrame({'filename':np.arange(len(os.listdir(path=root_dir)))})
for root, dirs, files in os.walk(root_dir):
for i, filename in enumerate(files):
self.df.loc[i, 'filename'] = filename
def __len__(self):
return len(os.listdir(path=self.root_dir))
def __getitem__(self, idx):
img_name = self.root_dir + self.df.loc[idx, 'filename']
image = Image.open(img_name)
image = torchvision.transforms.functional.to_tensor(image)
width = int(image.size(-1) / 2)
image, target = image[:, :, width:], image[:, :, :width]
if self.transform:
image, target = self.transform(image, target)
sample = {'x': image, 'target': target}
return sample