-
Notifications
You must be signed in to change notification settings - Fork 3
/
image_helpers.py
96 lines (70 loc) · 2.58 KB
/
image_helpers.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
"""
Note: the GitHub repo https://github.com/wasidennis/AdaptSegNet was used
as a starting point to for develop the code contained in this file.
"""
import os
import numpy as np
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
class ImageOps:
def __init__(self):
self.IMG_MEAN = np.array(
(104.00698793, 116.66876762, 122.67891434), dtype=np.float32)
self.LABELS_DICT = {0:"road", 1:"sidewalk", 2:"building", 3:"wall",
4:"fence", 5:"pole", 6:"light", 7:"sign",
8:"vegetation", 9:"terrain", 10:"sky", 11:"person",
12:"rider", 13:"car", 14:"truck", 15:"bus", 16:"train",
7:"motocycle", 18:"bicycle"
}
# 19x3 values, for Image's palette() module
self.PALETTE = [128, 64, 128, 244, 35, 232, 70, 70, 70, 102, 102, 156, 190,
153, 153, 153, 153, 153, 250, 170, 30, 220, 220, 0, 107, 142,
35, 152, 251, 152, 70, 130, 180, 220, 20, 60, 255, 0, 0, 0, 0,
142, 0, 0, 70, 0, 60, 100, 0, 80, 100, 0, 0, 230, 119, 11, 32]
zero_pad = 256 * 3 - len(self.PALETTE)
for i in range(zero_pad):
self.PALETTE.append(0)
def colorize_mask(self, mask):
# mask: numpy array of the mask
new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
new_mask.putpalette(self.PALETTE)
return new_mask
def get_concat_h(self, im1, im2):
dst = Image.new('RGB', (im1.width + im2.width, max(im1.height, im2.height)))
dst.paste(im1, (0, 0))
dst.paste(im2, (im1.width, 0))
return dst
def get_concat_v(self, im1, im2):
dst = Image.new('RGB', (max(im1.width, im2.width), im1.height + im2.height))
dst.paste(im1, (0, 0))
dst.paste(im2, (0, im1.height))
return dst
def process_image_for_saving(self, image, interp):
# handling RGB input image
image = interp(image).cpu().numpy().squeeze()
image = np.transpose(image, (1, 2, 0))
image += self.IMG_MEAN
image = image[:, :, ::-1]
image = image.astype(np.uint8)
image = Image.fromarray(image)
return image
def process_ade20k_image_for_saving(self, image, mean, std):
# handling RGB input image
image = image.cpu().numpy().squeeze()
image = np.transpose(image, (1, 2, 0))
image *= std
image += mean
image *= 255.
#image = image[:, :, ::-1]
image = image.astype(np.uint8)
return image
def save_concat_image(self, image, gt, pred, input_size, save_path, prefix):
"""
Save concatenation of image, ground truth and prediction
"""
image_concat = self.get_concat_v(image, pred)
image_concat = self.get_concat_v(gt, image_concat)
image_concat_path = os.path.join(
save_path, prefix+'_concat.png')
image_concat.save(image_concat_path)