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)