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) # =========================