-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathtransform.py
105 lines (76 loc) · 3.15 KB
/
transform.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
import numpy as np
import torch
from PIL import Image
def colormap_cityscapes(n):
cmap = np.zeros([n, 3]).astype(np.uint8)
cmap[0, :] = np.array([128, 64, 128])
cmap[1, :] = np.array([244, 35, 232])
cmap[2, :] = np.array([70, 70, 70])
cmap[3, :] = np.array([102, 102, 156])
cmap[4, :] = np.array([190, 153, 153])
cmap[5, :] = np.array([153, 153, 153])
cmap[6, :] = np.array([250, 170, 30])
cmap[7, :] = np.array([220, 220, 0])
cmap[8, :] = np.array([107, 142, 35])
cmap[9, :] = np.array([152, 251, 152])
cmap[10, :] = np.array([70, 130, 180])
cmap[11, :] = np.array([220, 20, 60])
cmap[12, :] = np.array([255, 0, 0])
cmap[13, :] = np.array([0, 0, 142])
cmap[14, :] = np.array([0, 0, 70])
cmap[15, :] = np.array([0, 60, 100])
cmap[16, :] = np.array([0, 80, 100])
cmap[17, :] = np.array([0, 0, 230])
cmap[18, :] = np.array([119, 11, 32])
cmap[19, :] = np.array([229, 23, 142]) # drivable-fallback
cmap[20, :] = np.array([156, 60, 200]) # non-drivable-fallback
cmap[21, :] = np.array([99, 250, 80]) # autorickshaw
cmap[22, :] = np.array([82, 92, 214]) # vehicle-fallback
cmap[23, :] = np.array([196, 209, 152]) # curb
cmap[24, :] = np.array([180, 165, 180]) # guard-rail
cmap[25, :] = np.array([37, 58, 77]) # billboard
cmap[26, :] = np.array([11, 35, 88]) # bs
cmap[27, :] = np.array([150, 100, 100]) # bridge
cmap[28, :] = np.array([255, 255, 255])
return cmap
def colormap(n):
cmap = np.zeros([n, 3]).astype(np.uint8)
for i in np.arange(n):
r, g, b = np.zeros(3)
for j in np.arange(8):
r = r + (1 << (7-j))*((i & (1 << (3*j))) >> (3*j))
g = g + (1 << (7-j))*((i & (1 << (3*j+1))) >> (3*j+1))
b = b + (1 << (7-j))*((i & (1 << (3*j+2))) >> (3*j+2))
cmap[i, :] = np.array([r, g, b])
return cmap
class Relabel:
def __init__(self, olabel, nlabel):
self.olabel = olabel
self.nlabel = nlabel
def __call__(self, tensor):
assert (isinstance(tensor, torch.LongTensor) or isinstance(
tensor, torch.ByteTensor)), 'tensor needs to be LongTensor'
tensor[tensor == self.olabel] = self.nlabel
return tensor
class ToLabel:
def __call__(self, image):
return torch.from_numpy(np.array(image)).long().unsqueeze(0)
class Colorize:
def __init__(self, n=22):
#self.cmap = colormap(256)
self.cmap = colormap_cityscapes(256)
self.cmap[n] = self.cmap[-1]
self.cmap = torch.from_numpy(self.cmap[:n])
def __call__(self, gray_image):
size = gray_image.size()
# print(size)
color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
#color_image = torch.ByteTensor(3, size[0], size[1]).fill_(0)
# for label in range(1, len(self.cmap)):
for label in range(0, len(self.cmap)):
mask = gray_image[0] == label
#mask = gray_image == label
color_image[0][mask] = self.cmap[label][0]
color_image[1][mask] = self.cmap[label][1]
color_image[2][mask] = self.cmap[label][2]
return color_image