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