-
Notifications
You must be signed in to change notification settings - Fork 139
/
Copy pathvis_utils.py
177 lines (151 loc) · 7.86 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
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
166
167
168
169
170
171
172
173
174
175
176
177
import numpy as np
import cv2
import torch
from PIL import Image
from importmagician import import_from
with import_from('./'):
from tools.culane_evaluation_py.culane_metric import culane_metric
from tools.tusimple_evaluation.lane import LaneEval
# Colors statics
BGR_RED = [0, 0, 255]
BGR_GREEN = [0, 255, 0]
BGR_BLUE = [255, 0, 0]
def tensor_image_to_numpy(images):
return (images * 255.0).cpu().numpy().astype(np.uint8)
def save_images(images, filenames):
# Save tensor images in range [0.0, 1.0]
# filenames: List[str]
assert images.shape[0] == len(filenames)
if type(images) != np.ndarray: # Flexible
images = tensor_image_to_numpy(images)
for i in range(len(filenames)):
Image.fromarray(images[i]).save(filenames[i])
def segmentation_visualize_batched(images, labels, colors, std=None, mean=None, trans=0.3, ignore_color=None,
auto_color=True, ignore_index=255):
# Draw images + labels from tensors (batched)
# images (4D), labels (3D), colors (2D), std, mean, ignore_color: torch.Tensor
# trans: how transparent is the label
# ignore_color: in range [0.0, 1.0]
assert images.shape[0] == labels.shape[0]
# Map label to RGB (N, d1, d2) = {0~20, 255} => (N, d1, d2, 3) = {0.0~1.0}
if colors is None: # Same color (white) for all classes
colors = torch.tensor([[0, 0, 0], [255, 255, 255]], device=images.device)
labels[labels > 0] = 1
else:
ignore_pixels = labels == ignore_index
bg_pixels = labels == 0
if auto_color: # Iterate colors (except background and ignore)
labels = (labels - 1) % (colors.shape[0] - 2) + 1
labels[ignore_pixels] = colors.shape[0] - 1 # Color for ignore
labels[bg_pixels] = 0
labels = colors[labels] / 255.0
# Denormalize if needed and map from (N, 3, d1, d2) to (N, d1, d2, 3)
images = images.permute(0, 2, 3, 1)
if std is not None and mean is not None:
images = (images.float() * std + mean).clamp_(0.0, 1.0)
# Mix (should not need another clamp)
results = images * trans + labels * (1 - trans)
if ignore_color is not None:
filter_mask = (labels == ignore_color).sum(dim=-1, keepdim=True) == ignore_color.shape[0]
results = results * ~filter_mask + images * filter_mask
return results
def draw_points(image, points, colors, radius=5, thickness=-1):
# Draw lines (defined by points) on an image as keypoints
# colors: can be a list that defines different colors for each line
for j in range(len(points)):
temp = points[j][(points[j][:, 0] > 0) * (points[j][:, 1] > 0)]
for k in range(temp.shape[0]):
color = colors[j] if isinstance(colors[0], list) else colors
cv2.circle(image, (int(temp[k][0]), int(temp[k][1])),
radius=radius, color=color, thickness=thickness)
return image
def draw_points_as_lines(image, points, colors, thickness=3):
# Draw lines (defined by points) on an image by connecting points to lines
# colors: can be a list that defines different colors for each line
for j in range(len(points)):
temp = points[j][(points[j][:, 0] > 0) * (points[j][:, 1] > 0)]
for k in range(temp.shape[0] - 1):
color = colors[j] if isinstance(colors[0], list) else colors
cv2.line(image,
(int(temp[k][0]), int(temp[k][1])),
(int(temp[k + 1][0]), int(temp[k + 1][1])),
color=color, thickness=thickness)
return image
def lane_detection_visualize_batched(images, masks=None, keypoints=None,
mask_colors=None, keypoint_color=None, std=None, mean=None,
control_points=None, gt_keypoints=None, style='point', line_trans=0.4,
compare_gt_metric='culane'):
# Draw images + lanes from tensors (batched)
# None masks/keypoints and keypoints (x < 0 or y < 0) will be ignored
# images (4D), masks (3D), keypoints (4D), colors (2D), std, mean: torch.Tensor
# keypoints can be either List[List[N x 2 numpy array]] (for variate length lanes) or a 4D numpy array
# filenames: List[str]
# keypoint_color: RGB
if masks is not None:
images = segmentation_visualize_batched(images, masks, mask_colors, std, mean,
trans=0, ignore_color=mask_colors[0])
if keypoints is not None:
if masks is None:
images = images.permute(0, 2, 3, 1)
if std is not None and mean is not None:
images = (images.float() * std + mean)
images = images.clamp_(0.0, 1.0) * 255.0
images = images[..., [2, 1, 0]].cpu().numpy().astype(np.uint8)
if keypoint_color is None:
keypoint_color = [0, 0, 0] # Black (sits well with lane colors)
else:
keypoint_color = keypoint_color[::-1] # To BGR
# Draw
for i in range(images.shape[0]):
# Compare with GT
if gt_keypoints is not None:
if compare_gt_metric == 'culane':
tp, fp, fn, pred_ious, _ = culane_metric(keypoints[i], gt_keypoints[i])
keypoint_color = [BGR_GREEN if iou >= 0.5 else BGR_RED for iou in pred_ious]
elif compare_gt_metric == 'tusimple':
x_pred = [keypoints[i][j][:, 0] for j in range(len(keypoints[i]))]
x_gt = [gt_keypoints[i][j][:, 0] for j in range(len(gt_keypoints[i]))]
y = gt_keypoints[i][0][:, 1].tolist()
acc, fp, fn, match, _ = LaneEval.bench_with_matches(x_pred, x_gt, y)
keypoint_color = [BGR_GREEN if m else BGR_RED for m in match]
if style == 'point':
if gt_keypoints is not None:
images[i] = draw_points(images[i], gt_keypoints[i], BGR_BLUE)
images[i] = draw_points(images[i], keypoints[i], keypoint_color)
elif style in ['line', 'bezier']:
overlay = images[i].copy()
if gt_keypoints is not None:
overlay = draw_points_as_lines(overlay, gt_keypoints[i], BGR_BLUE)
overlay = draw_points_as_lines(overlay, keypoints[i], keypoint_color)
images[i] = (images[i].astype(np.float) * line_trans +
overlay.astype(np.float) * (1 - line_trans)).astype(np.uint8)
if style == 'bezier':
assert control_points is not None, 'Must provide control points for style bezier!'
images[i] = draw_points(images[i], control_points[i], keypoint_color)
else:
raise ValueError('Unknown keypoint visualization style: {}\nPlease use point/line/bezier'.format(style))
images = images[..., [2, 1, 0]]
return images
def find_transform_by_name(cfg, name):
# Returns True if a transform name exists in augmentation cfg dict
if isinstance(cfg, dict):
if 'name' in cfg.keys() and cfg['name'] == name:
return True
else:
if 'transforms' in cfg.keys() and isinstance(cfg['transforms'], (list, tuple)):
return any([find_transform_by_name(t, name) for t in cfg['transforms']])
return False
def get_transform_attr_by_name(cfg, name, attr):
# Returns attr of first found transform by name in augmentation cfg dict
if isinstance(cfg, dict):
if 'name' in cfg.keys() and cfg['name'] == name:
return cfg.get(attr)
else:
if 'transforms' in cfg.keys() and isinstance(cfg['transforms'], (list, tuple)):
res = None
for t in cfg['transforms']:
res = get_transform_attr_by_name(t, name, attr)
if res is not None:
break
return res
return None