184 lines
6.7 KiB
Python
184 lines
6.7 KiB
Python
"""远程推理代理 — 摄像头子进程通过 Queue 向主进程 InferenceProcessPool 发起推理"""
|
||
import logging
|
||
import threading
|
||
import uuid
|
||
|
||
logger = logging.getLogger("analysis.remote_detector")
|
||
|
||
# 同一 resp_queue 只能有一个 drain 线程，否则多 RemoteDetector 会抢响应导致丢包
|
||
_drainers = {}
|
||
_drainers_lock = threading.Lock()
|
||
|
||
|
||
class _SharedResponseDrainer(object):
|
||
def __init__(self, resp_queue):
|
||
self._resp_q = resp_queue
|
||
self._pending = {}
|
||
self._lock = threading.Lock()
|
||
self._started = False
|
||
|
||
def ensure_started(self):
|
||
if self._started:
|
||
return
|
||
self._started = True
|
||
threading.Thread(
|
||
target=self._loop, name="remote-det-drain-%s" % id(self._resp_q), daemon=True,
|
||
).start()
|
||
|
||
def register(self, req_id):
|
||
evt = {"event": threading.Event(), "resp": None}
|
||
with self._lock:
|
||
self._pending[req_id] = evt
|
||
return evt
|
||
|
||
def unregister(self, req_id):
|
||
with self._lock:
|
||
self._pending.pop(req_id, None)
|
||
|
||
def _loop(self):
|
||
import queue as _q
|
||
while True:
|
||
try:
|
||
msg = self._resp_q.get(timeout=1.0)
|
||
except _q.Empty:
|
||
continue
|
||
if not msg:
|
||
continue
|
||
req_id = msg.get("req_id")
|
||
with self._lock:
|
||
item = self._pending.pop(req_id, None)
|
||
if item:
|
||
item["resp"] = msg
|
||
item["event"].set()
|
||
elif req_id:
|
||
logger.debug("remote_detector: 无匹配 pending req_id=%s", req_id)
|
||
|
||
|
||
def _get_drainer(resp_queue):
|
||
key = id(resp_queue)
|
||
with _drainers_lock:
|
||
drainer = _drainers.get(key)
|
||
if drainer is None:
|
||
drainer = _SharedResponseDrainer(resp_queue)
|
||
_drainers[key] = drainer
|
||
return drainer
|
||
|
||
|
||
class RemoteDetector(object):
|
||
ENGINE_NAME = "remote_pool"
|
||
|
||
def __init__(self, algorithm_spec, req_queue, resp_queue, timeout=35.0,
|
||
response_channel=None):
|
||
self._spec = algorithm_spec
|
||
self._req_q = req_queue
|
||
self._resp_q = resp_queue
|
||
self._timeout = timeout
|
||
self._response_channel = response_channel
|
||
self._drainer = _get_drainer(resp_queue)
|
||
|
||
def ready(self):
|
||
return self._req_q is not None and self._resp_q is not None
|
||
|
||
def load(self):
|
||
return True
|
||
|
||
def detect(self, frame):
|
||
if not self.ready():
|
||
raise RuntimeError("共享推理队列未就绪")
|
||
self._drainer.ensure_started()
|
||
try:
|
||
import cv2
|
||
orig_h, orig_w = frame.shape[:2]
|
||
inference_frame = self._maybe_downscale(frame)
|
||
infer_h, infer_w = inference_frame.shape[:2]
|
||
# 远处人员框很小，低质量 JPEG 会在 YOLO 前抹掉轮廓；95 在当前单路 CPU
|
||
# 场景仍可控，并显著缩小与直接单图推理的差异。
|
||
ok, buf = cv2.imencode(".jpg", inference_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
|
||
if not ok:
|
||
raise RuntimeError("推理帧编码失败")
|
||
jpeg = buf.tobytes()
|
||
except Exception as e:
|
||
raise RuntimeError("推理帧编码失败: %s" % e) from e
|
||
|
||
req_id = str(uuid.uuid4())
|
||
evt = self._drainer.register(req_id)
|
||
try:
|
||
self._req_q.put({
|
||
"req_id": req_id,
|
||
"algorithm": self._spec,
|
||
"jpeg": jpeg,
|
||
"response_channel": self._response_channel,
|
||
}, timeout=2.0)
|
||
except Exception as e:
|
||
self._drainer.unregister(req_id)
|
||
raise RuntimeError("推理请求入队失败: %s" % e) from e
|
||
|
||
if not evt["event"].wait(timeout=self._timeout):
|
||
self._drainer.unregister(req_id)
|
||
logger.warning("RemoteDetector 推理超时 algo=%s", self._spec.get("name"))
|
||
raise TimeoutError("共享推理响应超时（%ss）" % self._timeout)
|
||
resp = evt.get("resp") or {}
|
||
if not resp.get("ok"):
|
||
raise RuntimeError(resp.get("error") or "共享推理失败")
|
||
detections = resp.get("detections") or []
|
||
if infer_w != orig_w or infer_h != orig_h:
|
||
detections = self._restore_coordinates(
|
||
detections, float(orig_w) / infer_w, float(orig_h) / infer_h,
|
||
)
|
||
return detections
|
||
|
||
@staticmethod
|
||
def _restore_coordinates(detections, scale_x, scale_y):
|
||
"""把共享推理缩图坐标恢复到解码原帧坐标系。"""
|
||
result = []
|
||
for source in detections or []:
|
||
item = dict(source)
|
||
box = item.get("box")
|
||
if isinstance(box, (list, tuple)) and len(box) >= 4:
|
||
item["box"] = [float(box[0]) * scale_x, float(box[1]) * scale_y,
|
||
float(box[2]) * scale_x, float(box[3]) * scale_y]
|
||
keypoints = item.get("keypoints")
|
||
if isinstance(keypoints, list):
|
||
item["keypoints"] = [
|
||
[float(p[0]) * scale_x, float(p[1]) * scale_y] + list(p[2:])
|
||
if isinstance(p, (list, tuple)) and len(p) >= 2 else p
|
||
for p in keypoints
|
||
]
|
||
for polygon_key in ("polygon", "mask", "segments"):
|
||
polygon = item.get(polygon_key)
|
||
if isinstance(polygon, list):
|
||
item[polygon_key] = RemoteDetector._scale_polygon(polygon, scale_x, scale_y)
|
||
result.append(item)
|
||
return result
|
||
|
||
@staticmethod
|
||
def _scale_polygon(value, scale_x, scale_y):
|
||
scaled = []
|
||
for point in value:
|
||
if (isinstance(point, (list, tuple)) and len(point) >= 2 and
|
||
isinstance(point[0], (int, float)) and isinstance(point[1], (int, float))):
|
||
scaled.append([float(point[0]) * scale_x, float(point[1]) * scale_y] + list(point[2:]))
|
||
elif isinstance(point, list):
|
||
scaled.append(RemoteDetector._scale_polygon(point, scale_x, scale_y))
|
||
else:
|
||
scaled.append(point)
|
||
return scaled
|
||
|
||
def _maybe_downscale(self, frame):
|
||
try:
|
||
import cv2
|
||
h, w = frame.shape[:2]
|
||
max_side = max(
|
||
int(self._spec.get("input_width", 640) or 640),
|
||
int(self._spec.get("input_height", 640) or 640),
|
||
640,
|
||
) * 4
|
||
longest = max(h, w)
|
||
if longest <= max_side:
|
||
return frame
|
||
scale = float(max_side) / float(longest)
|
||
nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
|
||
return cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_AREA)
|
||
except Exception:
|
||
return frame
|