projects/human_detection/engine/threading_func.py

from queue import Queue
from threading import Event
from models.engine.threading_func import queue_clear

from projects.human_detection.engine.visualizer import Visualizer
import logging


def visualize_thread(visualizer_cfg, output_path, vis_img_batch_queue: Queue, track_queue: Queue, visualize_queue: Queue, eStop: Event, show_conf):
    logging.info("Start Visualize Thread")
    visualizer = Visualizer(**visualizer_cfg)
    input_video_info = vis_img_batch_queue.get()
    visualizer.init_writer(input_video_info, output_path)
    visualize_queue.put(input_video_info)
    track_item = track_queue.get()
    start_frame_idx = -1
    img_batch_item = vis_img_batch_queue.get()
    while (img_batch_item is not None and track_item is not None):
        if eStop.is_set():  break
        start_frame_idx, img_batch = img_batch_item
        track_start_frame_idx, track_result = track_item
        if (start_frame_idx != track_start_frame_idx) or (len(img_batch) != len(track_result)):
            error_msg=[501, f"Error when runing visualization 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
        for idx, (frame, frame_track_result) in enumerate(zip(img_batch, track_result)):

            visualizer.visualize(frame, frame_track_result, show_conf)
            visualize_queue.put(start_frame_idx + idx)
        img_batch_item = vis_img_batch_queue.get()
        track_item = track_queue.get()
    visualizer.close()

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