Spaces:
Runtime error
Runtime error
File size: 2,495 Bytes
a818fda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | from typing import List, Dict
import cv2
import numpy as np
from models.engine.visualizer import BaseVisualizer
class Visualizer(BaseVisualizer):
def __init__(self, fps: int=-1, min_width: int=-1):
""" Visualizer class for visualization (track_results + count_results).
Args:
class_map_ids (Dict): class mapping dictionary to map model's class to original class. Eg {0: 1, 1: 0, 2: 2, 3: 3} mean we swap class ID between 0 and 1.
fps (int): FPS for output video. If fps = -1, it will have same fps as input video.
min_width (int): minimum width for output video (height will be scaled to keep aspect ratio as input video). If min_width = -1, it will have same resolution as input video.
"""
class_names = ['pedestrian']
super().__init__(class_names, fps, min_width)
def visualize(self, img: np.ndarray, dettrack_at_frame_id: List[Dict]=None, show_conf: bool=True):
""" Function to visualize (track_results + count_results) a frame.
Args:
img (np.ndarray): image need to be visualized.
dettrack_at_frame_id (List[Dict]): batch of track results which can be obtained from Tracker class.
show_conf (bool): Visualize confidence of track results or not.
"""
# Draw tracking
if (dettrack_at_frame_id):
boxes = dettrack_at_frame_id["boxes"]
# classes = dettrack_at_frame_id["labels"]
ids = dettrack_at_frame_id["ids"]
for bbox, id_ in zip(boxes, ids):
id_ = int(id_)
score = bbox[4]
color = self.get_color(id_)
label = f'{id_}' + (f' {score:.2f}' if (show_conf) else '')
tl, tf = 2, 1
c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))
img = cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
img = cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)
img = cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
if (img.shape[0] != self.height or img.shape[1] != self.width):
img = cv2.resize(img, (self.width, self.height))
self.writer.write(img)
return img
|