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.")