L'électron rare
65bf3aad08
pose_filter _parse_env_* read POSE_FILTER* via VizConfig. multi_hmr_worker reads MULTIHMR_BACKEND/AUTOCAST/REMOTE via VizConfig. multihmr_remote reads JPEG/ASYNC/HOST/PORT via VizConfig. smplx_osc_sender reads AVBODY_HOST/REID/ALPHA via VizConfig. pose_bridge reads AVBODY_HOST/VDMX_* via VizConfig. iphone_usb_source reads CONCERT_MIRROR via VizConfig. lidar_calib reads ICP_LIDAR_EXTRINSIC via VizConfig. multihmr_coreml reads COREML_COMPUTE_UNITS via VizConfig.
465 lines
16 KiB
Python
465 lines
16 KiB
Python
"""Multi-HMR remote backend: drop-in replacement of MultiHMRCoreMLBackend
|
|
that delegates inference to a remote TCP server (see
|
|
``scripts/multihmr_server.py``).
|
|
|
|
Protocol (little-endian, persistent connection):
|
|
|
|
Request:
|
|
[4B uint32 payload_len]
|
|
[4B magic "REQ\x01"]
|
|
[1B uint8 format_id] # 1 = raw RGB uint8 HWC, 2 = JPEG (variable length)
|
|
[3B padding]
|
|
[variable image bytes] # IMG_BYTES if format=1, else JPEG bytes
|
|
[9 float32 K = 36 bytes]
|
|
|
|
The K block is *always* the last 36 bytes of the payload, regardless of
|
|
``format_id`` — the server slices it off before treating the rest as the
|
|
image.
|
|
|
|
Response:
|
|
[4B uint32 payload_len]
|
|
[4B magic "RSP\x01"]
|
|
[4B int32 status]
|
|
[v3d : 4*10475*3 f32]
|
|
[transl: 4*1*3 f32]
|
|
[scores: 4 f32]
|
|
[betas: 4*10 f32]
|
|
[expr : 4*10 f32]
|
|
|
|
Two extra features over the bare RPC:
|
|
|
|
* JPEG compression (``MULTIHMR_REMOTE_JPEG=1``, default ON, quality 80).
|
|
Cuts wire bytes from ~1.35 MB to ~50-150 KB.
|
|
|
|
* Asynchronous double-buffer (``MULTIHMR_REMOTE_ASYNC=1``, default ON).
|
|
``infer()`` is decoupled from the I/O round-trip via a dedicated worker
|
|
thread and two ``Queue(maxsize=1)`` slots. When the out-queue is empty
|
|
``infer()`` returns ``None`` — the worker loop reuses its last humans
|
|
list so the visualiser keeps drawing.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import queue
|
|
import socket
|
|
import struct
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
LOG = logging.getLogger("multihmr_remote")
|
|
|
|
IMG_SIZE = 672
|
|
N_PERSONS_FIXED = 4
|
|
N_VERTS = 10475
|
|
|
|
MAGIC_REQ = b"REQ\x01"
|
|
MAGIC_RSP = b"RSP\x01"
|
|
|
|
FORMAT_RAW = 1
|
|
FORMAT_JPEG = 2
|
|
|
|
IMG_BYTES = IMG_SIZE * IMG_SIZE * 3
|
|
K_BYTES = 9 * 4
|
|
REQ_HEADER = 4 + 1 + 3 # magic + format_id + 3-byte pad
|
|
# Fixed RAW-format request payload (mirrors RSP_PAYLOAD_LEN). The JPEG
|
|
# path is variable-length, so this is the upper-bound / RAW case only.
|
|
REQ_PAYLOAD_LEN = REQ_HEADER + IMG_BYTES + K_BYTES
|
|
|
|
V3D_BYTES = N_PERSONS_FIXED * N_VERTS * 3 * 4
|
|
TRANSL_BYTES = N_PERSONS_FIXED * 1 * 3 * 4
|
|
SCORES_BYTES = N_PERSONS_FIXED * 4
|
|
BETAS_BYTES = N_PERSONS_FIXED * 10 * 4
|
|
EXPR_BYTES = N_PERSONS_FIXED * 10 * 4
|
|
RSP_HEADER = 4 + 4
|
|
RSP_PAYLOAD_LEN = (RSP_HEADER + V3D_BYTES + TRANSL_BYTES
|
|
+ SCORES_BYTES + BETAS_BYTES + EXPR_BYTES)
|
|
|
|
|
|
def _env_flag(name: str, default: bool) -> bool:
|
|
raw = os.environ.get(name)
|
|
if raw is None:
|
|
return default
|
|
return raw.strip().lower() in ("1", "true", "yes", "on")
|
|
|
|
|
|
def _recv_exact(sock: socket.socket, n: int) -> bytes:
|
|
buf = bytearray(n)
|
|
view = memoryview(buf)
|
|
pos = 0
|
|
while pos < n:
|
|
got = sock.recv_into(view[pos:])
|
|
if got == 0:
|
|
raise ConnectionError("peer closed mid-stream")
|
|
pos += got
|
|
return bytes(buf)
|
|
|
|
|
|
def encode_request_raw(image_uint8_hwc: np.ndarray,
|
|
K_33: np.ndarray) -> bytes:
|
|
"""Raw uint8 HWC request (format_id=1, fixed payload length)."""
|
|
if image_uint8_hwc.shape != (IMG_SIZE, IMG_SIZE, 3):
|
|
raise ValueError(
|
|
f"image shape {image_uint8_hwc.shape} != "
|
|
f"({IMG_SIZE},{IMG_SIZE},3)")
|
|
if image_uint8_hwc.dtype != np.uint8:
|
|
raise ValueError(f"image dtype {image_uint8_hwc.dtype} != uint8")
|
|
K = np.ascontiguousarray(K_33, dtype="<f4").reshape(9)
|
|
img = np.ascontiguousarray(image_uint8_hwc, dtype=np.uint8)
|
|
img_bytes = img.tobytes()
|
|
header_after_magic = bytes([FORMAT_RAW, 0, 0, 0])
|
|
payload_len = REQ_HEADER + len(img_bytes) + K_BYTES
|
|
return b"".join([
|
|
struct.pack("<I", payload_len),
|
|
MAGIC_REQ,
|
|
header_after_magic,
|
|
img_bytes,
|
|
K.tobytes(),
|
|
])
|
|
|
|
|
|
def encode_request_jpeg(jpeg_bytes: bytes, K_33: np.ndarray) -> bytes:
|
|
"""JPEG request (format_id=2, variable payload length)."""
|
|
K = np.ascontiguousarray(K_33, dtype="<f4").reshape(9)
|
|
header_after_magic = bytes([FORMAT_JPEG, 0, 0, 0])
|
|
payload_len = REQ_HEADER + len(jpeg_bytes) + K_BYTES
|
|
return b"".join([
|
|
struct.pack("<I", payload_len),
|
|
MAGIC_REQ,
|
|
header_after_magic,
|
|
jpeg_bytes,
|
|
K.tobytes(),
|
|
])
|
|
|
|
|
|
def decode_response(payload: bytes) -> tuple[
|
|
np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
|
|
if len(payload) != RSP_PAYLOAD_LEN:
|
|
raise ValueError(
|
|
f"rsp payload {len(payload)} != {RSP_PAYLOAD_LEN}")
|
|
if payload[:4] != MAGIC_RSP:
|
|
raise ValueError(f"bad rsp magic {payload[:4]!r}")
|
|
status = struct.unpack("<i", payload[4:8])[0]
|
|
off = 8
|
|
v3d = np.frombuffer(payload, dtype="<f4",
|
|
count=N_PERSONS_FIXED * N_VERTS * 3,
|
|
offset=off).reshape(N_PERSONS_FIXED, N_VERTS, 3)
|
|
off += V3D_BYTES
|
|
transl = np.frombuffer(payload, dtype="<f4",
|
|
count=N_PERSONS_FIXED * 1 * 3,
|
|
offset=off).reshape(N_PERSONS_FIXED, 1, 3)
|
|
off += TRANSL_BYTES
|
|
scores = np.frombuffer(payload, dtype="<f4",
|
|
count=N_PERSONS_FIXED,
|
|
offset=off).reshape(N_PERSONS_FIXED)
|
|
off += SCORES_BYTES
|
|
betas = np.frombuffer(payload, dtype="<f4",
|
|
count=N_PERSONS_FIXED * 10,
|
|
offset=off).reshape(N_PERSONS_FIXED, 10)
|
|
off += BETAS_BYTES
|
|
expr = np.frombuffer(payload, dtype="<f4",
|
|
count=N_PERSONS_FIXED * 10,
|
|
offset=off).reshape(N_PERSONS_FIXED, 10)
|
|
return (v3d.copy(), transl.copy(), scores.copy(),
|
|
betas.copy(), expr.copy(), int(status))
|
|
|
|
|
|
# Back-compat shim — old call sites used encode_request(img, K) for raw.
|
|
def encode_request(image_uint8_hwc: np.ndarray, K_33: np.ndarray) -> bytes:
|
|
return encode_request_raw(image_uint8_hwc, K_33)
|
|
|
|
|
|
class _Tensorlike:
|
|
"""Mimics CoreMLArray to avoid a hard import on multihmr_coreml."""
|
|
__slots__ = ("_arr",)
|
|
|
|
def __init__(self, arr: np.ndarray) -> None:
|
|
self._arr = arr
|
|
|
|
def detach(self) -> "_Tensorlike":
|
|
return self
|
|
|
|
def cpu(self) -> "_Tensorlike":
|
|
return self
|
|
|
|
def numpy(self) -> np.ndarray:
|
|
return self._arr
|
|
|
|
def item(self) -> float:
|
|
return float(self._arr.reshape(-1)[0])
|
|
|
|
@property
|
|
def shape(self) -> tuple[int, ...]:
|
|
return tuple(self._arr.shape)
|
|
|
|
|
|
def _humans_from_arrays(v3d: np.ndarray, transl: np.ndarray,
|
|
scores: np.ndarray, betas: np.ndarray,
|
|
expr: np.ndarray, det_thresh: float
|
|
) -> list[dict[str, Any]]:
|
|
humans: list[dict[str, Any]] = []
|
|
for k in range(N_PERSONS_FIXED):
|
|
sc = float(scores[k])
|
|
if sc < det_thresh:
|
|
continue
|
|
humans.append({
|
|
"v3d": _Tensorlike(v3d[k]),
|
|
"transl_pelvis": _Tensorlike(transl[k]),
|
|
"scores": _Tensorlike(np.array([sc], dtype=np.float32)),
|
|
"shape": _Tensorlike(betas[k]),
|
|
"expression": _Tensorlike(expr[k]),
|
|
})
|
|
return humans
|
|
|
|
|
|
class MultiHMRRemoteBackend:
|
|
"""TCP client backend mirroring ``MultiHMRCoreMLBackend.infer`` API.
|
|
|
|
JPEG compression and async double-buffering are toggleable via env
|
|
(``MULTIHMR_REMOTE_JPEG``, ``MULTIHMR_REMOTE_ASYNC``).
|
|
"""
|
|
|
|
def __init__(self, host: str = "192.168.0.175", port: int = 57140,
|
|
connect_timeout: float = 3.0,
|
|
io_timeout: float = 5.0) -> None:
|
|
self.host = host
|
|
self.port = port
|
|
self.connect_timeout = connect_timeout
|
|
self.io_timeout = io_timeout
|
|
self._sock: socket.socket | None = None
|
|
self._lock = threading.Lock()
|
|
|
|
from .config import VizConfig as _VizConfig
|
|
_cfg = _VizConfig.from_env()
|
|
self.use_jpeg = _cfg.multihmr_remote_jpeg
|
|
self.jpeg_quality = _cfg.multihmr_remote_jpeg_quality
|
|
self.use_async = _cfg.multihmr_remote_async
|
|
|
|
# Async pipeline state.
|
|
# Multi-buffer queues (2 in / 3 out) absorb jitter without
|
|
# stalling capture. Drop-oldest semantics on overflow.
|
|
self._in_q: queue.Queue[tuple[bytes, float, float]] = queue.Queue(
|
|
maxsize=2)
|
|
self._out_q: queue.Queue[
|
|
tuple[list[dict[str, Any]], dict[str, float]]
|
|
] = queue.Queue(maxsize=3)
|
|
self._stop = threading.Event()
|
|
self._async_det_thresh = 0.3
|
|
self._worker_thread: threading.Thread | None = None
|
|
self._last_stats: dict[str, float] = {}
|
|
|
|
if self.use_jpeg:
|
|
try:
|
|
import cv2 # noqa: F401
|
|
except ImportError:
|
|
LOG.warning("cv2 unavailable client-side, disabling JPEG")
|
|
self.use_jpeg = False
|
|
|
|
if self.use_async:
|
|
self._start_worker()
|
|
LOG.info(
|
|
"MultiHMRRemoteBackend %s:%d (jpeg=%s q=%d, async=%s)",
|
|
host, port, self.use_jpeg, self.jpeg_quality, self.use_async)
|
|
|
|
# -- connection management -------------------------------------------
|
|
|
|
def _connect(self) -> socket.socket:
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.settimeout(self.connect_timeout)
|
|
sock.connect((self.host, self.port))
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
sock.settimeout(self.io_timeout)
|
|
LOG.info("connected to %s:%d", self.host, self.port)
|
|
return sock
|
|
|
|
def _ensure_sock(self) -> socket.socket:
|
|
if self._sock is None:
|
|
self._sock = self._connect()
|
|
return self._sock
|
|
|
|
def _drop_sock(self) -> None:
|
|
if self._sock is not None:
|
|
try:
|
|
self._sock.close()
|
|
except OSError:
|
|
pass
|
|
self._sock = None
|
|
|
|
@staticmethod
|
|
def is_available(host: str | None = None, port: int | None = None
|
|
) -> bool:
|
|
from .config import VizConfig as _VizConfig
|
|
_cfg = _VizConfig.from_env()
|
|
host = host or _cfg.multihmr_remote_host
|
|
port = port or _cfg.multihmr_remote_port
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(1.0)
|
|
s.connect((host, port))
|
|
s.close()
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
# -- request encoding ------------------------------------------------
|
|
|
|
def _encode_request_from_chw(
|
|
self, image_chw_float32: np.ndarray, K_33: np.ndarray
|
|
) -> tuple[bytes, float]:
|
|
"""Return (request bytes, encode_ms)."""
|
|
img = np.asarray(image_chw_float32, dtype=np.float32)
|
|
if img.ndim == 4 and img.shape[0] == 1:
|
|
img = img[0]
|
|
if img.shape != (3, IMG_SIZE, IMG_SIZE):
|
|
raise ValueError(
|
|
f"image shape {img.shape} != (3,{IMG_SIZE},{IMG_SIZE})")
|
|
img_hwc = np.clip(img.transpose(1, 2, 0) * 255.0, 0.0, 255.0
|
|
).astype(np.uint8)
|
|
K = np.asarray(K_33, dtype=np.float32)
|
|
if K.ndim == 3 and K.shape[0] == 1:
|
|
K = K[0]
|
|
if K.shape != (3, 3):
|
|
raise ValueError(f"K shape {K.shape} != (3,3)")
|
|
|
|
t0 = time.monotonic()
|
|
if self.use_jpeg:
|
|
import cv2 # local import to keep optional
|
|
# cv2.imencode wants BGR for nicest JPEG perceptually but the
|
|
# server decodes back to RGB ; encode RGB->BGR once for parity.
|
|
bgr = cv2.cvtColor(img_hwc, cv2.COLOR_RGB2BGR)
|
|
ok, enc = cv2.imencode(
|
|
".jpg", bgr,
|
|
[int(cv2.IMWRITE_JPEG_QUALITY), self.jpeg_quality])
|
|
if not ok:
|
|
raise RuntimeError("cv2.imencode failed")
|
|
req = encode_request_jpeg(bytes(enc), K)
|
|
else:
|
|
req = encode_request_raw(img_hwc, K)
|
|
enc_ms = (time.monotonic() - t0) * 1e3
|
|
return req, enc_ms
|
|
|
|
# -- synchronous fallback -------------------------------------------
|
|
|
|
def _send_recv(self, req: bytes) -> bytes:
|
|
attempts = 0
|
|
last_err: Exception | None = None
|
|
while attempts < 2:
|
|
attempts += 1
|
|
try:
|
|
sock = self._ensure_sock()
|
|
sock.sendall(req)
|
|
len_buf = _recv_exact(sock, 4)
|
|
payload_len = struct.unpack("<I", len_buf)[0]
|
|
if payload_len != RSP_PAYLOAD_LEN:
|
|
raise ValueError(
|
|
f"unexpected rsp len {payload_len}")
|
|
return _recv_exact(sock, payload_len)
|
|
except (ConnectionError, BrokenPipeError, OSError,
|
|
socket.timeout) as e:
|
|
LOG.warning("rpc failed (try %d): %s", attempts, e)
|
|
self._drop_sock()
|
|
last_err = e
|
|
raise RuntimeError(f"remote inference failed: {last_err}")
|
|
|
|
# -- async worker ---------------------------------------------------
|
|
|
|
def _start_worker(self) -> None:
|
|
self._worker_thread = threading.Thread(
|
|
target=self._async_loop, name="multihmr-remote",
|
|
daemon=True)
|
|
self._worker_thread.start()
|
|
|
|
def _async_loop(self) -> None:
|
|
while not self._stop.is_set():
|
|
try:
|
|
req, t_submit, det_thresh = self._in_q.get(timeout=0.5)
|
|
except queue.Empty:
|
|
continue
|
|
t_send = time.monotonic()
|
|
try:
|
|
with self._lock:
|
|
payload = self._send_recv(req)
|
|
except Exception as e: # noqa: BLE001
|
|
LOG.warning("async send_recv failed: %s", e)
|
|
continue
|
|
t_recv = time.monotonic()
|
|
try:
|
|
v3d, transl, scores, betas, expr, status = decode_response(
|
|
payload)
|
|
except Exception as e: # noqa: BLE001
|
|
LOG.warning("decode_response failed: %s", e)
|
|
continue
|
|
if status != 0:
|
|
humans: list[dict[str, Any]] = []
|
|
else:
|
|
humans = _humans_from_arrays(
|
|
v3d, transl, scores, betas, expr, det_thresh)
|
|
stats = {
|
|
"queue_wait_ms": (t_send - t_submit) * 1e3,
|
|
"rpc_ms": (t_recv - t_send) * 1e3,
|
|
}
|
|
# Drop any pending stale output before pushing.
|
|
try:
|
|
self._out_q.get_nowait()
|
|
except queue.Empty:
|
|
pass
|
|
try:
|
|
self._out_q.put_nowait((humans, stats))
|
|
except queue.Full:
|
|
pass
|
|
|
|
# -- public API -----------------------------------------------------
|
|
|
|
def infer(
|
|
self,
|
|
image_chw_float32: np.ndarray,
|
|
K_33: np.ndarray,
|
|
det_thresh: float = 0.3,
|
|
) -> list[dict[str, Any]] | None:
|
|
"""In sync mode returns the humans list (possibly empty).
|
|
|
|
In async mode, submits the new frame (non-blocking, drop-newest
|
|
if previous frame still in flight) and returns whatever output
|
|
is ready in the out-queue. Returns ``None`` if nothing is ready
|
|
yet — caller must reuse its last humans list.
|
|
"""
|
|
req, _enc_ms = self._encode_request_from_chw(
|
|
image_chw_float32, K_33)
|
|
|
|
if not self.use_async:
|
|
with self._lock:
|
|
payload = self._send_recv(req)
|
|
v3d, transl, scores, betas, expr, status = decode_response(
|
|
payload)
|
|
if status != 0:
|
|
return []
|
|
return _humans_from_arrays(
|
|
v3d, transl, scores, betas, expr, det_thresh)
|
|
|
|
# Async path.
|
|
self._async_det_thresh = det_thresh
|
|
# drop-newest semantics: keep the freshest pending frame
|
|
try:
|
|
self._in_q.get_nowait()
|
|
except queue.Empty:
|
|
pass
|
|
try:
|
|
self._in_q.put_nowait((req, time.monotonic(), det_thresh))
|
|
except queue.Full:
|
|
pass
|
|
|
|
try:
|
|
humans, stats = self._out_q.get_nowait()
|
|
except queue.Empty:
|
|
return None
|
|
self._last_stats = stats
|
|
return humans
|
|
|
|
def close(self) -> None:
|
|
self._stop.set()
|
|
with self._lock:
|
|
self._drop_sock()
|