|
|
| 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
|
|
|
