models/engine/visualizer.py

from abc import abstractmethod
from typing import List, Optional
import cv2
import subprocess
import numpy as np 

def putText(img, text: str, position, 
               text_font: int=0, text_scale: int=1,
               bg_color=(255,255,255),
               text_color=(255,0,255),
               bg_thickness=8, 
               text_thickness=1, 
               lineType=cv2.LINE_AA):
    """ Function to put text on image.
    
    Args:
        img (_type_): 
        text (str): _description_
        position (_type_): Top-left position of text.
        text_font (int, optional): font size of text. Defaults to 0.
        text_scale (int, optional): text scale. Defaults to 1.
        bg_color (tuple, optional): text background color. Defaults to (255,255,255).
        text_color (tuple, optional): text foreground color. Defaults to (255,0,255).
        bg_thickness (int, optional): text background thickness. Defaults to 8.
        text_thickness (int, optional): text foreground thickness. Defaults to 1.
        lineType (_type_, optional): line type. Defaults to cv2.LINE_AA.

    Returns:
        _type_: _description_
    """
    img = cv2.putText(img, text, position, text_font, text_scale, bg_color, thickness=bg_thickness, lineType=lineType)
    img = cv2.putText(img, text, position, text_font, text_scale, text_color, thickness=text_thickness, lineType=lineType)
    return img 


class BaseVisualizer():

    def __init__(self, class_names: Optional[List[str]], 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.
        """

        self.fps = fps
        self.min_width = min_width
        self.class_names = class_names
    
    def init_writer(self, input_video_info: List[int], output_path: str):
        """ Init video writer for write visualized frame to output video.

        Args:
            input_video_info (List[int]): It is a list that includes 4 elements of input video information (fps, width, height, num_frames).
            output_path (str): Path to save output video. 
        """
        if (self.fps == -1):
            self.fps = input_video_info[0]
        self.width, self.height = input_video_info[1], input_video_info[2]
        if (self.min_width > 0):
            out_width = min(self.min_width, self.width)
            self.height = (self.height * out_width)//self.width
            self.width = out_width
        self.output_path = output_path
        self.writer = cv2.VideoWriter(self.output_path, cv2.VideoWriter_fourcc(*"mp4v"), int(self.fps), (self.width, self.height))
        
    @staticmethod
    def get_color(idx):
        idx = idx * 3
        color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)

        return color

    @staticmethod
    def draw_dash_line(img,pt1,pt2,color,thickness=1,style='dotted',gap=20):
        dist =((pt1[0]-pt2[0])**2+(pt1[1]-pt2[1])**2)**.5
        pts= []
        for i in  np.arange(0,dist,gap):
            r=i/dist
            x=int((pt1[0]*(1-r)+pt2[0]*r)+.5)
            y=int((pt1[1]*(1-r)+pt2[1]*r)+.5)
            p = (x,y)
            pts.append(p)
        if len(pts) ==0: 
            return
        if style=='dotted':
            for p in pts:
                cv2.circle(img,p,thickness,color,-1)
        else:
            s=pts[0]
            e=pts[0]
            i=0
            for p in pts:
                s=e
                e=p
                if i%2==1:
                    cv2.line(img,s,e,color,thickness)
                i+=1

    @staticmethod
    def draw_dash_poly(img,pts,color,thickness=1,style='dotted',gap=20):
        """ draw a polygon with dash line.

        Args:
            img (_type_): input image.
            pts (_type_): _description_
            color (_type_): _description_
            thickness (int, optional): _description_. Defaults to 1.
            style (str, optional): _description_. Defaults to 'dotted'.
            gap (int, optional): _description_. Defaults to 20.

        Returns:
            _type_: _description_
        """
        s=pts[0]
        e=pts[0]
        pts.append(pts.pop(0))
        for p in pts:
            s=e
            e=p
            BaseVisualizer.draw_dash_line(img,s,e,color,thickness,style,gap=gap)
        return img

    @staticmethod
    def draw_dash_rect(img,pt1,pt2,color,thickness=1,style='dotted',gap=10):
        pts = [pt1,(pt2[0],pt1[1]),pt2,(pt1[0],pt2[1])] 
        return BaseVisualizer.draw_dash_poly(img,pts,color,thickness,style,gap=gap)   
    
    def close(self):
        """ Function to release video writer. It should be called after finish visualization for all input frames.
        """
        self.writer.release()
    
    def convert(self):
        subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {self.output_path} -c:v libx264 {self.output_path}".split())

    @abstractmethod
    def visualize(self, *args,**kwargs):
        """ Each project should implement this function to visualize a frame.

        """
        raise NotImplementedError