Spaces:
Runtime error
Runtime error
File size: 8,748 Bytes
a818fda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | 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.") |