Skip to content
Snippets Groups Projects
Commit ebf59c89 authored by karolmajek's avatar karolmajek Committed by Kai Chen
Browse files

Per class colors in Mask results viz (#1834)


* per class color in mask rcnn

* fix max label id

* pass flake8

* correct the annotation

Co-authored-by: default avatarWenwei Zhang <40779233+ZwwWayne@users.noreply.github.com>
Co-authored-by: default avatarKai Chen <chenkaidev@gmail.com>
parent dcfedf54
No related branches found
No related tags found
No related merge requests found
......@@ -148,20 +148,26 @@ def show_result(img,
else:
bbox_result, segm_result = result, None
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
# draw segmentation masks
if segm_result is not None:
segms = mmcv.concat_list(segm_result)
inds = np.where(bboxes[:, -1] > score_thr)[0]
np.random.seed(42)
color_masks = [
np.random.randint(0, 256, (1, 3), dtype=np.uint8)
for _ in range(max(labels) + 1)
]
for i in inds:
color_mask = np.random.randint(0, 256, (1, 3), dtype=np.uint8)
i = int(i)
color_mask = color_masks[labels[i]]
mask = maskUtils.decode(segms[i]).astype(np.bool)
img[mask] = img[mask] * 0.5 + color_mask * 0.5
# draw bounding boxes
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
mmcv.imshow_det_bboxes(
img,
bboxes,
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment