Spaces:
Runtime error
Runtime error
| 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 | |