|
|
|
from .base import RectangleVisualizer, TextVisualizer |
|
|
|
|
|
class BoundingBoxVisualizer: |
|
def __init__(self): |
|
self.rectangle_visualizer = RectangleVisualizer() |
|
|
|
def visualize(self, image_bgr, boxes_xywh): |
|
for bbox_xywh in boxes_xywh: |
|
image_bgr = self.rectangle_visualizer.visualize(image_bgr, bbox_xywh) |
|
return image_bgr |
|
|
|
|
|
class ScoredBoundingBoxVisualizer: |
|
def __init__(self, bbox_visualizer_params=None, score_visualizer_params=None, **kwargs): |
|
if bbox_visualizer_params is None: |
|
bbox_visualizer_params = {} |
|
if score_visualizer_params is None: |
|
score_visualizer_params = {} |
|
self.visualizer_bbox = RectangleVisualizer(**bbox_visualizer_params) |
|
self.visualizer_score = TextVisualizer(**score_visualizer_params) |
|
|
|
def visualize(self, image_bgr, scored_bboxes): |
|
boxes_xywh, box_scores = scored_bboxes |
|
assert len(boxes_xywh) == len( |
|
box_scores |
|
), "Number of bounding boxes {} should be equal to the number of scores {}".format( |
|
len(boxes_xywh), len(box_scores) |
|
) |
|
for i, box_xywh in enumerate(boxes_xywh): |
|
score_i = box_scores[i] |
|
image_bgr = self.visualizer_bbox.visualize(image_bgr, box_xywh) |
|
score_txt = "{0:6.4f}".format(score_i) |
|
topleft_xy = box_xywh[0], box_xywh[1] |
|
image_bgr = self.visualizer_score.visualize(image_bgr, score_txt, topleft_xy) |
|
return image_bgr |
|
|
