未验证 提交 28197297 编写于 作者: K Kaipeng Deng 提交者: GitHub

refine bbox visualize (#2594)

* refine bbox visualize

* use different color for different categories

* change color range
上级 91bc8c5b
...@@ -48,9 +48,9 @@ def draw_mask(image, im_id, segms, threshold, alpha=0.7): ...@@ -48,9 +48,9 @@ def draw_mask(image, im_id, segms, threshold, alpha=0.7):
""" """
Draw mask on image Draw mask on image
""" """
im_width, im_height = image.size
mask_color_id = 0 mask_color_id = 0
w_ratio = .4 w_ratio = .4
color_list = colormap(rgb=True)
img_array = np.array(image).astype('float32') img_array = np.array(image).astype('float32')
for dt in np.array(segms): for dt in np.array(segms):
if im_id != dt['image_id']: if im_id != dt['image_id']:
...@@ -59,7 +59,6 @@ def draw_mask(image, im_id, segms, threshold, alpha=0.7): ...@@ -59,7 +59,6 @@ def draw_mask(image, im_id, segms, threshold, alpha=0.7):
if score < threshold: if score < threshold:
continue continue
mask = mask_util.decode(segm) * 255 mask = mask_util.decode(segm) * 255
color_list = colormap(rgb=True)
color_mask = color_list[mask_color_id % len(color_list), 0:3] color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1 mask_color_id += 1
for c in range(3): for c in range(3):
...@@ -77,34 +76,43 @@ def draw_bbox(image, im_id, catid2name, bboxes, threshold, ...@@ -77,34 +76,43 @@ def draw_bbox(image, im_id, catid2name, bboxes, threshold,
""" """
draw = ImageDraw.Draw(image) draw = ImageDraw.Draw(image)
catid2color = {}
color_list = colormap(rgb=True)[:40]
for dt in np.array(bboxes): for dt in np.array(bboxes):
if im_id != dt['image_id']: if im_id != dt['image_id']:
continue continue
catid, bbox, score = dt['category_id'], dt['bbox'], dt['score'] catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
if score < threshold: if score < threshold:
continue continue
xmin, ymin, w, h = bbox
xmin, ymin, w, h = bbox
if is_bbox_normalized: if is_bbox_normalized:
im_width, im_height = image.size im_width, im_height = image.size
xmin *= im_width xmin *= im_width
ymin *= im_height ymin *= im_height
w *= im_width w *= im_width
h *= im_height h *= im_height
xmax = xmin + w xmax = xmin + w
ymax = ymin + h ymax = ymin + h
if catid not in catid2color:
idx = np.random.randint(len(color_list))
catid2color[catid] = color_list[idx]
color = tuple(catid2color[catid])
# draw bbox
draw.line( draw.line(
[(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
(xmin, ymin)], (xmin, ymin)],
width=2, width=2,
fill='red') fill=color)
if image.mode == 'RGB':
text = catid2name[catid] # draw label
tw, th = draw.textsize(text) text = "{} {:.2f}".format(catid2name[catid], score)
draw.rectangle([(xmin + 1, ymin + 1), tw, th = draw.textsize(text)
(xmin + tw + 1, ymin + th + 1)], draw.rectangle([(xmin + 1, ymin - th),
fill='red') (xmin + tw + 1, ymin)],
draw.text((xmin + 1, ymin + 1), text, fill=(255, 255, 255)) fill=color)
draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
return image return image
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册