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