models/engine/threading_func.py

from queue import Queue, Full, Empty
from threading import Event
from mmcv import VideoReader
import logging
from gradio import Progress
from models.trackers.byte_track import BYTETracker
import torch
import numpy as np

def queue_clear(q: Queue):
    """ Clear all items in the queue.

    Args:
        q (Queue): input queue.
    """
    with q.mutex: q.queue.clear()

def queue_get(q: Queue, eStop: Event, retry_interval=1, item_idx=None, default_item=None):
    """wrapper for queue.get() with timeout, retry and event stop.

    Args:
        q (Queue): input queue.
        eStop (Event): event to stop the thread.
        retry_interval (int, optional): time to wait before retry to get the item. Defaults to 1 second.
        item_idx (_type_, optional): index of item to get. This is used for logging information. Defaults to None.
        default_item (_type_, optional): default item to return if error or early stop. Defaults to None.

    Returns:
        any: item in the queue.
    """
    if not q.empty():
        return q.get()

    while not eStop.is_set():
        try:
            item = q.get(timeout=retry_interval)
            return item
        except Empty:
            if item_idx is not None:
                logging.info(f"Waiting to get item {item_idx}")
    if item_idx is not None:
        logging.info(f"Early Stop. Return Default item at iter {item_idx}")
    return default_item

def queue_put(q: Queue, item, eStop: Event, retry_interval=1, item_idx=None):
    """ wrapper for queue.put() with timeout, retry and event stop.

    Args:
        q (Queue): input queue.
        item (_type_): item to put in the queue.
        eStop (Event): event to stop the thread.
        retry_interval (int, optional): time to wait before retry to put the item. Defaults to 1 second.
        item_idx (_type_, optional): index of item to put. This is used for logging information. Defaults to None.
    """
    if not q.full():
        q.put(item)
        return 

    while not eStop.is_set():
        try:
            q.put(item, timeout=retry_interval)
            return
        except Full:
            if item_idx is not None:
                logging.info(f"Waiting to put item at {item_idx}")
    if item_idx is not None:
        logging.info(f"Early Stop. No item is put at iter {item_idx}")  

def batch_extract_thread(video_path: str, 
                         img_batch_queue: Queue, 
                         vis_img_batch_queue: Queue, 
                         eStop: Event, 
                         batch_size=32):
    """Thread function to extract a batch of frames from video and put it to img_batch_queue and vis_img_batch_queue.

    Args:
        video_path (str): input video path.
        img_batch_queue (Queue): output queue for batch of frames, used for processing.
        vis_img_batch_queue (Queue): output queue for batch of frames, used for visualization.
        eStop (Event): event to stop the thread.
        batch_size (int, optional): number of images in a batch. Defaults to 32.
    """
    logging.info("Start Batch Extract Thread")
    vidcap = VideoReader(video_path)
    vis_img_batch_queue.put([vidcap.fps, vidcap.width, vidcap.height, len(vidcap)])
    start_frame_idx = 0
    last_frame_idx = len(vidcap)
    end_frame_idx = start_frame_idx

    while (start_frame_idx < last_frame_idx):
        if eStop.is_set(): break
        end_frame_idx = min(start_frame_idx + batch_size, last_frame_idx)
        img_batch = []
        for frame_idx in range(start_frame_idx, end_frame_idx):
            img = vidcap[frame_idx]
            if (img is None):
                break
            img_batch.append(img)
        if (len(img_batch) == 0):
            break
        item_data = [start_frame_idx, img_batch]
        queue_put(img_batch_queue, item_data, eStop)
        queue_put(vis_img_batch_queue, item_data , eStop)
        start_frame_idx = end_frame_idx
    
    if eStop.is_set():
        queue_clear(img_batch_queue)
        queue_clear(vis_img_batch_queue)
    else:
        logging.info(f"Finish batch_extract_thread for video_file {video_path} at end_frame_idx {end_frame_idx}.")
    img_batch_queue.put(None)
    vis_img_batch_queue.put(None)

def detect_thread(obj_detector, 
                  img_batch_queue: Queue, 
                  det_queue: Queue, 
                  eStop: Event,
                  put_img_batch: bool=False):
    """ detect_thread function to run detection on a batch of frames.

    Args:
        obj_detector (_type_): object detector, for example YOLOV7TRT/-ONXX.
        img_batch_queue (Queue): input queue for batch of frames, which is the output from batch_extract_thread.
        det_queue (Queue): output queue for detection results.
        eStop (Event): event to stop the thread.
        put_img_batch (bool, optional): If True, the input image batch will also put tp det_queue. 
                This is often used for later step that require images of detected objects, such Human Pose or ReID.
                Defaults to False.
    """
    logging.info("Start Detection Thread")
    item = img_batch_queue.get()
    start_frame_idx = -1
    while item is not None:
        if eStop.is_set(): break
        start_frame_idx, img_batch = item
        logging.info(f"Run detection at frame idx: {start_frame_idx}")
        try:
            det_result = obj_detector.infer_batch(img_batch)
            item_data = [start_frame_idx, det_result]
            if (put_img_batch):
                item_data.append(img_batch)
            queue_put(det_queue, item_data, eStop)
        except Exception as e:
            error_msg=[501, f"Error when running detection at frame idx: {start_frame_idx}]. "]
            log_error_message = f"{error_msg[1]}. Error {e}"
            logging.exception(log_error_message)
            eStop.set()
            break  
        item = img_batch_queue.get()  

    # Finish this thread. 
    if eStop.is_set():
        logging.warning(f"Early stop detect_thread at start_frame_idx {start_frame_idx}")
        queue_clear(det_queue)
    else:
        logging.info(f"Finish detect_thread at start_frame_idx {start_frame_idx}.")    
    det_queue.put(None)

def bytetrack_thread(tracker_cfg, det_queue: Queue, track_queue: Queue, eStop: Event, conf_thres: float):
    logging.info("Start Tracking Thread")
    tracker = BYTETracker(
        **tracker_cfg
    )
    item = det_queue.get()
    start_frame_idx = -1
    
    while item is not None:
        if eStop.is_set():break
        start_frame_idx, det_result = item

        if isinstance(det_result[0]['boxes'],np.ndarray):
            det_result = [{key:torch.from_numpy(value) for key,value in dict_det.items()} for dict_det in det_result]
            
        try:
            track_result = tracker.track_batch(start_frame_idx,det_result,conf_thres)
        except Exception as e:
            error_msg=[501,f"Error when running tracking at start_frame_idx {start_frame_idx}: {e}"]
            log_error_message = f"Error {error_msg[0]}: {error_msg[1]}"
            logging.error(log_error_message)
            eStop.set()
            break
        queue_put(track_queue, [start_frame_idx, track_result], eStop)
        item = det_queue.get()

    # Finish this thread 
    if eStop.is_set():
        logging.warning(f"Early stop at start_frame_idx {start_frame_idx}.")
        queue_clear(track_queue)
    else:
        logging.info(f"Finish track_thread.")
    track_queue.put(None)
    
def update_progress_thread(visualize_queue: Queue, progress: Progress, eStop: Event):
    """Show the progress of the video processing on Gradio, measured by the number of frames visualized.

    Args:
        visualize_queue (Queue): input queue for batch of frames, which is the output from batch_extract_thread.
        progress (Progress): Gradio progress bar.
        eStop (Event): event to stop the thread.
    """
    
    fps, width, height, total_num_frames = visualize_queue.get()
    progress(0, desc="Starting...")
    start_frame_idx = -1
    for frame_idx in progress.tqdm(range(total_num_frames), total=total_num_frames):
        item = visualize_queue.get()
        if (item is None):
            break
        start_frame_idx = item
        if (start_frame_idx != frame_idx):
            error_msg=[501, f"Error when runing update progress at start_frame_idx {start_frame_idx}. "]
            log_error_message = f"Error {error_msg[0]}: {error_msg[1]}"
            logging.error(log_error_message)
            eStop.set()
            break
    
    # Finish this thread    
    if eStop.is_set():
        logging.warning(f"Early stop at start_frame_idx {start_frame_idx}")
    else:
        logging.info(f"Finish update_progress_thread.")