-
Notifications
You must be signed in to change notification settings - Fork 12
/
vis_utils.py
61 lines (56 loc) · 2.07 KB
/
vis_utils.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
import cv2
from PIL import Image
import numpy as np
import copy
import torchvision.transforms as transforms
class ImgLoader(object):
def __init__(self, img_size: int):
self.img_size = img_size
self.transform = transforms.Compose([
transforms.Resize((510, 510), Image.BILINEAR),
transforms.CenterCrop((img_size, img_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])
def load(self, image_path: str):
ori_img = cv2.imread(image_path)
assert ori_img.shape[2] == 3, "3(RGB) channels is required."
img = copy.deepcopy(ori_img)
img = img[:, :, ::-1] # convert BGR to RGB
img = Image.fromarray(img)
img = self.transform(img)
# center crop
ori_img = cv2.resize(ori_img, (510, 510))
pad = (510 - self.img_size) // 2
ori_img = ori_img[pad:pad+self.img_size, pad:pad+self.img_size]
return img, ori_img
def get_cdict():
_jet_data = {
# 'red': ((0.00, 0, 0),
# (0.35, 0, 0),
# (0.66, 1, 1),
# (0.89, 1, 1),
# (1.00, 0.5, 0.5)),
'red': ((0.00, 0, 0),
(0.35, 0.5, 0.5),
(0.66, 1, 1),
(0.89, 1, 1),
(1.00, 0.8, 0.8)),
'green': ((0.000, 0, 0),
(0.125, 0, 0),
(0.375, 1, 1),
(0.640, 1, 1),
(0.910, 0.3, 0.3),
(1.000, 0, 0)),
# 'blue': ((0.00, 0.5, 0.5),
# (0.11, 1, 1),
# (0.34, 1, 1),
# (0.65, 0, 0),
# (1.00, 0, 0))}
'blue': ((0.00, 0.30, 0.30),
(0.25, 0.8, 0.8),
(0.34, 0.8, 0.8),
(0.65, 0, 0),
(1.00, 0, 0))
}
return _jet_data