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.
248 lines
9.2 KiB
Python
248 lines
9.2 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
|
|
|
|
INFERENCE_DEVICE = DEVICE if torch.cuda.is_available() and int(DEVICE) >= 0 else "cpu"
|
|
INFERENCE_SIZE_FULL = IMG_SIZE_FULL if torch.cuda.is_available() else min(IMG_SIZE_FULL, 640)
|
|
INFERENCE_HALF = bool(USE_HALF and torch.cuda.is_available())
|
|
|
|
# Async YOLO worker
|
|
# =========================
|
|
|
|
|
|
def raw_yolo_boxes(result, offset_x=0, offset_y=0, scale=1.0, pad_x=0.0, pad_y=0.0):
|
|
if result.boxes is None or len(result.boxes) == 0:
|
|
return []
|
|
xyxy = result.boxes.xyxy.detach().cpu().numpy()
|
|
confs = result.boxes.conf.detach().cpu().numpy()
|
|
clss = result.boxes.cls.detach().cpu().numpy().astype(int)
|
|
out = []
|
|
for b, c, cls_id in zip(xyxy, confs, clss):
|
|
x1, y1, x2, y2 = map(float, b)
|
|
out.append(np.array([
|
|
(x1 - pad_x) * scale + offset_x,
|
|
(y1 - pad_y) * scale + offset_y,
|
|
(x2 - pad_x) * scale + offset_x,
|
|
(y2 - pad_y) * scale + offset_y,
|
|
float(c),
|
|
float(cls_id),
|
|
], dtype=np.float32))
|
|
return out
|
|
|
|
|
|
def fixed_letterbox(image, size):
|
|
h, w = image.shape[:2]
|
|
scale = min(float(size) / max(1, w), float(size) / max(1, h))
|
|
new_w = max(1, min(int(size), int(round(w * scale))))
|
|
new_h = max(1, min(int(size), int(round(h * scale))))
|
|
interpolation = cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LINEAR
|
|
resized = cv2.resize(image, (new_w, new_h), interpolation=interpolation)
|
|
left = (int(size) - new_w) // 2
|
|
top = (int(size) - new_h) // 2
|
|
canvas = np.full((int(size), int(size), 3), 114, dtype=np.uint8)
|
|
canvas[top:top + new_h, left:left + new_w] = resized
|
|
return canvas, scale, left, top
|
|
|
|
|
|
class YOLOWorker:
|
|
def __init__(self, model, full_frame_shape=None):
|
|
self.model = model
|
|
self.full_frame_shape = full_frame_shape
|
|
self.req = deque(maxlen=YOLO_QUEUE_MAX)
|
|
self.res = deque(maxlen=1)
|
|
self.lock = threading.Lock()
|
|
self.request_event = threading.Event()
|
|
self.ready_event = threading.Event()
|
|
self.startup_error = None
|
|
self.last_error = None
|
|
self.running = False
|
|
self.thread = threading.Thread(target=self._loop, daemon=True)
|
|
|
|
def start(self):
|
|
self.running = True
|
|
self.thread.start()
|
|
if not self.ready_event.wait(timeout=120.0):
|
|
self.stop()
|
|
raise RuntimeError("YOLO worker startup timed out")
|
|
if self.startup_error is not None:
|
|
raise RuntimeError(f"YOLO worker startup failed: {self.startup_error}") from self.startup_error
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
with self.lock:
|
|
self.req.clear()
|
|
self.request_event.set()
|
|
self.thread.join(timeout=5.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))
|
|
self.request_event.set()
|
|
|
|
def try_get(self):
|
|
with self.lock:
|
|
if not self.res:
|
|
return None
|
|
return self.res.pop()
|
|
|
|
def _loop(self):
|
|
try:
|
|
self._warmup()
|
|
except Exception as exc:
|
|
self.startup_error = exc
|
|
self.ready_event.set()
|
|
return
|
|
self.ready_event.set()
|
|
|
|
while self.running:
|
|
self.request_event.wait(timeout=0.2)
|
|
self.request_event.clear()
|
|
item = None
|
|
with self.lock:
|
|
if self.req:
|
|
item = self.req.pop()
|
|
self.req.clear()
|
|
|
|
if item is None:
|
|
continue
|
|
|
|
frame, roi_box, mode, ts = item
|
|
h, w = frame.shape[:2]
|
|
frame_infer = preprocess_for_yolo(frame)
|
|
dets = []
|
|
raw_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
|
|
crop_model, input_scale, pad_x, pad_y = fixed_letterbox(crop, IMG_SIZE_ROI)
|
|
t0 = time.perf_counter()
|
|
with torch.inference_mode():
|
|
r = self.model(
|
|
crop_model,
|
|
conf=YOLO_CONF_EFFECTIVE,
|
|
imgsz=IMG_SIZE_ROI,
|
|
verbose=False,
|
|
max_det=MAX_DET,
|
|
device=INFERENCE_DEVICE,
|
|
half=INFERENCE_HALF
|
|
)[0]
|
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
|
raw_dets = raw_yolo_boxes(
|
|
r,
|
|
offset_x=ox,
|
|
offset_y=oy,
|
|
scale=1.0 / input_scale,
|
|
pad_x=pad_x,
|
|
pad_y=pad_y,
|
|
)
|
|
dets = filter_yolo_boxes_with_scores(
|
|
r,
|
|
frame_w=w,
|
|
frame_h=h,
|
|
offset_x=ox,
|
|
offset_y=oy,
|
|
min_conf=BT_LOW,
|
|
input_scale=input_scale,
|
|
pad_x=pad_x,
|
|
pad_y=pad_y,
|
|
content_w=crop.shape[1],
|
|
content_h=crop.shape[0],
|
|
)
|
|
|
|
else:
|
|
sh, sw = frame_infer.shape[:2]
|
|
short = min(sh, sw)
|
|
scale = 1.0
|
|
target = INFERENCE_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=INFERENCE_SIZE_FULL,
|
|
verbose=False,
|
|
max_det=MAX_DET,
|
|
device=INFERENCE_DEVICE,
|
|
half=INFERENCE_HALF
|
|
)[0]
|
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
|
raw_small = raw_yolo_boxes(r)
|
|
|
|
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
|
|
raw_dets = [
|
|
np.array([d[0] * inv, d[1] * inv, d[2] * inv, d[3] * inv, d[4], d[5]], dtype=np.float32)
|
|
for d in raw_small
|
|
]
|
|
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:
|
|
raw_dets = raw_small
|
|
dets = dets_s
|
|
|
|
except Exception as exc:
|
|
dets = []
|
|
raw_dets = []
|
|
infer_ms = 0.0
|
|
message = f"{type(exc).__name__}: {exc}"
|
|
if message != self.last_error:
|
|
print(f"[yolo] inference failed: {message}", flush=True)
|
|
self.last_error = message
|
|
|
|
with self.lock:
|
|
self.res.append((dets, ts, mode, infer_ms, used_roi, raw_dets))
|
|
|
|
def _warmup(self):
|
|
if not torch.cuda.is_available():
|
|
return
|
|
dummy_roi = np.zeros((IMG_SIZE_ROI, IMG_SIZE_ROI, 3), dtype=np.uint8)
|
|
full_h, full_w = self.full_frame_shape or (IMG_SIZE_FULL, IMG_SIZE_FULL)
|
|
dummy_full = np.zeros((max(1, int(full_h)), max(1, int(full_w)), 3), dtype=np.uint8)
|
|
with torch.inference_mode():
|
|
for frame, size in ((dummy_roi, IMG_SIZE_ROI), (dummy_full, IMG_SIZE_FULL)):
|
|
self.model(
|
|
frame,
|
|
conf=YOLO_CONF_EFFECTIVE,
|
|
imgsz=size,
|
|
verbose=False,
|
|
max_det=MAX_DET,
|
|
device=INFERENCE_DEVICE,
|
|
half=USE_HALF,
|
|
)
|
|
print(f"Model warmed up in YOLO worker at {IMG_SIZE_ROI} and {IMG_SIZE_FULL}", flush=True)
|
|
|
|
|
|
# =========================
|