def _draw_bbox(self, img, bbox):
bbox = bbox[0]
left_top = (bbox[0], bbox[1])
right_bottom = (bbox[2], bbox[3])
cv2.rectangle(img,
left_top,
right_bottom,
color_val('red'),
thickness=self.thickness)
cv2.putText(img, 'no wear helmet', (int(bbox[0]), int(bbox[1]) - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, color_val('red'), 2)
return img
def _draw_bbox(self, img, bbox):
bbox = bbox[0]
left_top = (bbox[0], bbox[1])
right_bottom = (bbox[2], bbox[3])
cv2.rectangle(img,
left_top,
right_bottom,
color_val('red'),
thickness=self.thickness)
return img
def imshow_det_bboxes(img,
bboxes,
labels,
class_names=None,
score_thr=0,
bbox_color='green',
text_color='green',
thickness=1,
font_scale=0.5,
show=True,
win_name='',
wait_time=0,
out_file=None,
color_dist=None):
"""Draw bboxes and class labels (with scores) on an image.
Args:
img (str or ndarray): The image to be displayed.
bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or
(n, 5).
labels (ndarray): Labels of bboxes.
class_names (list[str]): Names of each classes.
score_thr (float): Minimum score of bboxes to be shown.
bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
text_color (str or tuple or :obj:`Color`): Color of texts.
thickness (int): Thickness of lines.
font_scale (float): Font scales of texts.
show (bool): Whether to show the image.
win_name (str): The window name.
wait_time (int): Value of waitKey param.
out_file (str or None): The filename to write the image.
color_dist(dist):the bbox color dist
"""
assert bboxes.ndim == 2
assert labels.ndim == 1
assert bboxes.shape[0] == labels.shape[0]
assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
img = mmcv.imread(img)
if score_thr > 0:
assert bboxes.shape[1] == 5
scores = bboxes[:, -1]
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds]
bbox_color = color_val(bbox_color)
text_color = color_val(text_color)
for bbox, label in zip(bboxes, labels):
bbox_int = bbox.astype(np.int32)
left_top = (bbox_int[0], bbox_int[1])
right_bottom = (bbox_int[2], bbox_int[3])
label_text = class_names[
label] if class_names is not None else 'cls {}'.format(label)
bbox_color = color_val(color_dist[label_text])
text_color = color_val(color_dist[label_text])
cv2.rectangle(img,
left_top,
right_bottom,
bbox_color,
thickness=thickness)
if len(bbox) > 4:
label_text += '|{:.02f}'.format(bbox[-1])
cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2),
cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color)
def imshow_tracklets(img,
bboxes,
labels=None,
ids=None,
thickness=2,
font_scale=0.4,
show=False,
win_name='',
color=None,
out_file=None):
assert bboxes.ndim == 2
assert labels.ndim == 1
assert bboxes.shape[0] == labels.shape[0]
# assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
if isinstance(img, str):
img = imread(img)
i = 0
if bboxes.shape[0] == 0:
if out_file is not None:
imwrite(img, out_file)
return img
if isinstance(bboxes, torch.Tensor):
bboxes = bboxes.numpy()
labels = labels.numpy()
ids = ids.numpy()
for bbox, label in zip(bboxes, labels):
x1, y1, x2, y2, _ = bbox.astype(np.int32)
if ids is not None:
if color is None:
bbox_color = random_color(ids[i])
bbox_color = [int(255 * _c) for _c in bbox_color][::-1]
else:
bbox_color = mmcv.color_val(color)
img[y1:y1 + 12, x1:x1 + 20, :] = bbox_color
cv2.putText(img,
str(ids[i]), (x1, y1 + 10),
cv2.FONT_HERSHEY_COMPLEX,
font_scale,
color=color_val('black'))
else:
if color is None:
bbox_color = color_val('green')
else:
bbox_color = mmcv.color_val(color)
cv2.rectangle(img, (x1, y1), (x2, y2), bbox_color, thickness=thickness)
if bbox[-1] < 0:
bbox[-1] = np.nan
# label_text = '{:.02f}'.format(bbox[-1])
# img[y1 - 12:y1, x1:x1 + 30, :] = bbox_color
# cv2.putText(
# img,
# label_text, (x1, y1 - 2),
# cv2.FONT_HERSHEY_COMPLEX,
# font_scale,
# color=color_val('black'))
i += 1
if show:
imshow(img, win_name)
if out_file is not None:
imwrite(img, out_file)
return img
imgs = data.get_data(images)
print("start model recongnize")
with torch.no_grad():
output = model(return_loss=False, **imgs)
outputs.append(output)
use_softmax = True
if use_softmax is True:
print("Averaging score over {} clips with softmax".format(
outputs[0].shape[0]))
results = [softmax(res, dim=1).mean(axis=0) for res in outputs]
else:
print("Averaging score over {} clips without softmax (ie, raw)".format(
outputs[0].shape[0]))
results = [res.mean(axis=0) for res in outputs]
pred = int(np.argmax(results, axis=1))
print(pred)
print(results[0][pred])
print(gt_labels[pred])
text_color = color_val('green')
for image in images:
cv2.putText(image, str(gt_labels[pred]), (100, 100),
cv2.FONT_HERSHEY_COMPLEX, 0.5, text_color)
cv2.imshow('image', image)
cv2.waitKey(0)
cap.release()
cv2.destroyAllWindows()