diff --git a/PaddleCV/object_detection/ppdet/utils/visualizer.py b/PaddleCV/object_detection/ppdet/utils/visualizer.py index c1983a62a4e57eecbb1ba4ca04565770b5f57fc5..21f787a84ee548da698e706c6dc06acc34684168 100644 --- a/PaddleCV/object_detection/ppdet/utils/visualizer.py +++ b/PaddleCV/object_detection/ppdet/utils/visualizer.py @@ -48,9 +48,9 @@ def draw_mask(image, im_id, segms, threshold, alpha=0.7): """ Draw mask on image """ - im_width, im_height = image.size mask_color_id = 0 w_ratio = .4 + color_list = colormap(rgb=True) img_array = np.array(image).astype('float32') for dt in np.array(segms): if im_id != dt['image_id']: @@ -59,7 +59,6 @@ def draw_mask(image, im_id, segms, threshold, alpha=0.7): if score < threshold: continue mask = mask_util.decode(segm) * 255 - color_list = colormap(rgb=True) color_mask = color_list[mask_color_id % len(color_list), 0:3] mask_color_id += 1 for c in range(3): @@ -77,34 +76,43 @@ def draw_bbox(image, im_id, catid2name, bboxes, threshold, """ draw = ImageDraw.Draw(image) + catid2color = {} + color_list = colormap(rgb=True)[:40] for dt in np.array(bboxes): if im_id != dt['image_id']: continue catid, bbox, score = dt['category_id'], dt['bbox'], dt['score'] if score < threshold: continue - xmin, ymin, w, h = bbox + xmin, ymin, w, h = bbox if is_bbox_normalized: im_width, im_height = image.size xmin *= im_width ymin *= im_height w *= im_width h *= im_height - xmax = xmin + w 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( [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)], width=2, - fill='red') - if image.mode == 'RGB': - text = catid2name[catid] - tw, th = draw.textsize(text) - draw.rectangle([(xmin + 1, ymin + 1), - (xmin + tw + 1, ymin + th + 1)], - fill='red') - draw.text((xmin + 1, ymin + 1), text, fill=(255, 255, 255)) + fill=color) + + # draw label + text = "{} {:.2f}".format(catid2name[catid], score) + tw, th = draw.textsize(text) + draw.rectangle([(xmin + 1, ymin - th), + (xmin + tw + 1, ymin)], + fill=color) + draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255)) return image