You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
141 lines
4.6 KiB
Python
141 lines
4.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
import time
|
|
import threading
|
|
from collections import deque
|
|
import torch
|
|
|
|
from config import *
|
|
from helpers import clip_box, crop_roi, preprocess_for_yolo, filter_yolo_boxes_with_scores
|
|
|
|
# Async YOLO worker
|
|
# =========================
|
|
|
|
class YOLOWorker:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.req = deque(maxlen=YOLO_QUEUE_MAX)
|
|
self.res = deque(maxlen=1)
|
|
self.lock = threading.Lock()
|
|
self.running = False
|
|
self.thread = threading.Thread(target=self._loop, daemon=True)
|
|
|
|
def start(self):
|
|
self.running = True
|
|
self.thread.start()
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
self.thread.join(timeout=1.0)
|
|
|
|
def submit(self, frame_eff_bgr, roi_box_eff, mode, ts):
|
|
with self.lock:
|
|
self.req.append((frame_eff_bgr, roi_box_eff, mode, ts))
|
|
|
|
def try_get(self):
|
|
with self.lock:
|
|
if not self.res:
|
|
return None
|
|
return self.res.pop()
|
|
|
|
def _loop(self):
|
|
while self.running:
|
|
item = None
|
|
with self.lock:
|
|
if self.req:
|
|
item = self.req.pop()
|
|
self.req.clear()
|
|
|
|
if item is None:
|
|
time.sleep(0.001)
|
|
continue
|
|
|
|
frame, roi_box, mode, ts = item
|
|
h, w = frame.shape[:2]
|
|
frame_infer = preprocess_for_yolo(frame)
|
|
dets = []
|
|
infer_ms = 0.0
|
|
used_roi = False
|
|
|
|
try:
|
|
if roi_box is not None:
|
|
roi_box = clip_box(roi_box, w, h)
|
|
crop, ox, oy = crop_roi(frame_infer, roi_box)
|
|
if crop.size > 0:
|
|
used_roi = True
|
|
t0 = time.perf_counter()
|
|
with torch.inference_mode():
|
|
r = self.model(
|
|
crop,
|
|
conf=YOLO_CONF_EFFECTIVE,
|
|
imgsz=IMG_SIZE_ROI,
|
|
verbose=False,
|
|
max_det=MAX_DET,
|
|
device=DEVICE,
|
|
half=USE_HALF
|
|
)[0]
|
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
|
dets = filter_yolo_boxes_with_scores(
|
|
r,
|
|
frame_w=w,
|
|
frame_h=h,
|
|
offset_x=ox,
|
|
offset_y=oy,
|
|
min_conf=BT_LOW
|
|
)
|
|
|
|
else:
|
|
sh, sw = frame_infer.shape[:2]
|
|
short = min(sh, sw)
|
|
scale = 1.0
|
|
target = IMG_SIZE_FULL
|
|
if short > target:
|
|
scale = target / float(short)
|
|
small = cv2.resize(
|
|
frame_infer,
|
|
(int(sw * scale), int(sh * scale)),
|
|
interpolation=cv2.INTER_AREA
|
|
)
|
|
else:
|
|
small = frame_infer
|
|
|
|
t0 = time.perf_counter()
|
|
with torch.inference_mode():
|
|
r = self.model(
|
|
small,
|
|
conf=YOLO_CONF_EFFECTIVE,
|
|
imgsz=IMG_SIZE_FULL,
|
|
verbose=False,
|
|
max_det=MAX_DET,
|
|
device=DEVICE,
|
|
half=USE_HALF
|
|
)[0]
|
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
|
|
|
dets_s = filter_yolo_boxes_with_scores(
|
|
r,
|
|
frame_w=small.shape[1],
|
|
frame_h=small.shape[0],
|
|
offset_x=0,
|
|
offset_y=0,
|
|
min_conf=BT_LOW
|
|
)
|
|
if scale != 1.0:
|
|
inv = 1.0 / scale
|
|
dets = [
|
|
np.array([d[0] * inv, d[1] * inv, d[2] * inv, d[3] * inv, d[4]], dtype=np.float32)
|
|
for d in dets_s
|
|
]
|
|
else:
|
|
dets = dets_s
|
|
|
|
except Exception:
|
|
dets = []
|
|
infer_ms = 0.0
|
|
|
|
with self.lock:
|
|
self.res.append((dets, ts, mode, infer_ms, used_roi))
|
|
|
|
|
|
# =========================
|