/
nyuv2_loader.py
165 lines (135 loc) · 5.16 KB
/
nyuv2_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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import os
import collections
import torch
import numpy as np
import scipy.misc as m
from torch.utils import data
from ptsemseg.utils import recursive_glob
from ptsemseg.augmentations import Compose, RandomHorizontallyFlip, RandomRotate, Scale
class NYUv2Loader(data.Dataset):
"""
NYUv2 loader
Download From (only 13 classes):
test source: http://www.doc.ic.ac.uk/~ahanda/nyu_test_rgb.tgz
train source: http://www.doc.ic.ac.uk/~ahanda/nyu_train_rgb.tgz
test_labels source:
https://github.com/ankurhanda/nyuv2-meta-data/raw/master/test_labels_13/nyuv2_test_class13.tgz
train_labels source:
https://github.com/ankurhanda/nyuv2-meta-data/raw/master/train_labels_13/nyuv2_train_class13.tgz
"""
def __init__(
self,
root,
split="training",
is_transform=False,
img_size=(480, 640),
augmentations=None,
img_norm=True,
test_mode=False
):
self.root = root
self.is_transform = is_transform
self.n_classes = 14
self.augmentations = augmentations
self.img_norm = img_norm
self.test_mode = test_mode
self.img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
self.mean = np.array([104.00699, 116.66877, 122.67892])
self.files = collections.defaultdict(list)
self.cmap = self.color_map(normalized=False)
split_map = {"training": "train", "val": "test"}
self.split = split_map[split]
for split in ["train", "test"]:
file_list = recursive_glob(rootdir=self.root + split + "/", suffix="png")
self.files[split] = file_list
def __len__(self):
return len(self.files[self.split])
def __getitem__(self, index):
img_path = self.files[self.split][index].rstrip()
img_number = img_path.split("_")[-1][:4]
lbl_path = os.path.join(
self.root, self.split + "_annot", "new_nyu_class13_" + img_number + ".png"
)
img = m.imread(img_path)
img = np.array(img, dtype=np.uint8)
lbl = m.imread(lbl_path)
lbl = np.array(lbl, dtype=np.uint8)
if not (len(img.shape) == 3 and len(lbl.shape) == 2):
return self.__getitem__(np.random.randint(0, self.__len__()))
if self.augmentations is not None:
img, lbl = self.augmentations(img, lbl)
if self.is_transform:
img, lbl = self.transform(img, lbl)
return img, lbl
def transform(self, img, lbl):
img = m.imresize(img, (self.img_size[0], self.img_size[1])) # uint8 with RGB mode
img = img[:, :, ::-1] # RGB -> BGR
img = img.astype(np.float64)
img -= self.mean
if self.img_norm:
# Resize scales images from 0 to 255, thus we need
# to divide by 255.0
img = img.astype(float) / 255.0
# NHWC -> NCHW
img = img.transpose(2, 0, 1)
classes = np.unique(lbl)
lbl = lbl.astype(float)
lbl = m.imresize(lbl, (self.img_size[0], self.img_size[1]), "nearest", mode="F")
lbl = lbl.astype(int)
assert np.all(classes == np.unique(lbl))
img = torch.from_numpy(img).float()
lbl = torch.from_numpy(lbl).long()
return img, lbl
def color_map(self, N=256, normalized=False):
"""
Return Color Map in PASCAL VOC format
"""
def bitget(byteval, idx):
return (byteval & (1 << idx)) != 0
dtype = "float32" if normalized else "uint8"
cmap = np.zeros((N, 3), dtype=dtype)
for i in range(N):
r = g = b = 0
c = i
for j in range(8):
r = r | (bitget(c, 0) << 7 - j)
g = g | (bitget(c, 1) << 7 - j)
b = b | (bitget(c, 2) << 7 - j)
c = c >> 3
cmap[i] = np.array([r, g, b])
cmap = cmap / 255.0 if normalized else cmap
return cmap
def decode_segmap(self, temp):
r = temp.copy()
g = temp.copy()
b = temp.copy()
for l in range(0, self.n_classes):
r[temp == l] = self.cmap[l, 0]
g[temp == l] = self.cmap[l, 1]
b[temp == l] = self.cmap[l, 2]
rgb = np.zeros((temp.shape[0], temp.shape[1], 3))
rgb[:, :, 0] = r / 255.0
rgb[:, :, 1] = g / 255.0
rgb[:, :, 2] = b / 255.0
return rgb
if __name__ == "__main__":
import matplotlib.pyplot as plt
augmentations = Compose([Scale(512), RandomRotate(10), RandomHorizontallyFlip()])
local_path = "/home/meet/datasets/NYUv2/"
dst = NYUv2Loader(local_path, is_transform=True, augmentations=augmentations)
bs = 4
trainloader = data.DataLoader(dst, batch_size=bs, num_workers=0)
for i, datas in enumerate(trainloader):
imgs, labels = datas
imgs = imgs.numpy()[:, ::-1, :, :]
imgs = np.transpose(imgs, [0, 2, 3, 1])
f, axarr = plt.subplots(bs, 2)
for j in range(bs):
axarr[j][0].imshow(imgs[j])
axarr[j][1].imshow(dst.decode_segmap(labels.numpy()[j]))
plt.show()
a = input()
if a == "ex":
break
else:
plt.close()