L'électron rare
f58f1d40e8
Add FAIL_THRESHOLD=30 counter: after 30 consecutive inference failures (NotImplementedError or any Exception), log once at ERROR and exit the loop instead of spinning at full CPU. Reset on success.
213 lines
7.2 KiB
Python
213 lines
7.2 KiB
Python
"""Worker NLF : capture webcam Mac, inference TorchScript multi-personne,
|
|
extraction vertices 3D nonparametriques SMPL (6890 verts), ecriture State.
|
|
|
|
NLF (Sarandi, NeurIPS 2024) fournit des vertices directement via le path
|
|
nonparametrique — pas besoin de modele SMPL-X externe. Le checkpoint
|
|
TorchScript est auto-contenu : detecteur + estimateur multi-personne.
|
|
|
|
Cadence cible : 8-12 fps sur M5 (NLF-L). NLF-S pour > 15 fps.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from .state import NLFPerson, State
|
|
|
|
LOG = logging.getLogger("nlf")
|
|
|
|
CACHE = Path.home() / ".cache" / "av-live-nlf"
|
|
CKPT_L = CACHE / "nlf_l_multi.torchscript"
|
|
CKPT_S = CACHE / "nlf_s_multi.torchscript"
|
|
|
|
N_VERTS = 6890
|
|
N_JOINTS = 24
|
|
|
|
FAIL_THRESHOLD = 30 # ~1 s at 30 fps before giving up
|
|
|
|
|
|
class NLFWorker:
|
|
def __init__(self, state: State, num_persons: int = 4,
|
|
target_fps: float = 10.0, device: str = "mps",
|
|
use_small: bool = False) -> None:
|
|
self.state = state
|
|
self.num_persons = num_persons
|
|
self.period = 1.0 / max(1.0, target_fps)
|
|
self.device = device
|
|
self.ckpt_path = CKPT_S if use_small else CKPT_L
|
|
self._stop = threading.Event()
|
|
self._thread: threading.Thread | None = None
|
|
self._smooth_pos: list[list] = []
|
|
self.failure_count = 0
|
|
|
|
@staticmethod
|
|
def is_available() -> bool:
|
|
return CKPT_L.exists() or CKPT_S.exists()
|
|
|
|
def start(self) -> None:
|
|
self._thread = threading.Thread(
|
|
target=self._run, name="nlf", daemon=True)
|
|
self._thread.start()
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
|
|
def _record_success(self) -> None:
|
|
self.failure_count = 0
|
|
|
|
def _run(self) -> None:
|
|
try:
|
|
import torch
|
|
import cv2
|
|
except ImportError as e:
|
|
LOG.error("deps manquantes : %s — uv sync --extra nlf", e)
|
|
return
|
|
|
|
if self.device == "mps" and not torch.backends.mps.is_available():
|
|
LOG.warning("MPS unavailable, falling back to cpu")
|
|
device = "cpu"
|
|
else:
|
|
device = self.device
|
|
|
|
if not self.ckpt_path.exists():
|
|
if CKPT_L.exists():
|
|
self.ckpt_path = CKPT_L
|
|
elif CKPT_S.exists():
|
|
self.ckpt_path = CKPT_S
|
|
else:
|
|
LOG.error("No NLF checkpoint found in %s", CACHE)
|
|
return
|
|
|
|
try:
|
|
model = torch.jit.load(
|
|
str(self.ckpt_path), map_location=device).eval()
|
|
except Exception as e:
|
|
LOG.error("NLF load failed: %s", e)
|
|
return
|
|
ckpt_name = self.ckpt_path.stem
|
|
LOG.info("NLF loaded (%s) on %s", ckpt_name, device)
|
|
|
|
from .euro_filter import OneEuroFilter
|
|
from .tracker import IoUTracker
|
|
from .state import PoseKp
|
|
self._smooth_pos = [
|
|
[OneEuroFilter(0.8, 0.05) for _ in range(3)]
|
|
for _ in range(self.num_persons)
|
|
]
|
|
tracker = IoUTracker(iou_threshold=0.20, max_miss=8)
|
|
|
|
cap = cv2.VideoCapture(0)
|
|
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
|
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
|
if not cap.isOpened():
|
|
LOG.error("camera index 0 indisponible")
|
|
return
|
|
LOG.info("camera ouverte")
|
|
|
|
while not self._stop.is_set():
|
|
t0 = time.monotonic()
|
|
ok, frame_bgr = cap.read()
|
|
if not ok:
|
|
time.sleep(self.period)
|
|
continue
|
|
|
|
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
|
tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1)
|
|
frame_batch = tensor.unsqueeze(0).to(device)
|
|
|
|
try:
|
|
with torch.inference_mode():
|
|
pred = model.detect_smpl_batched(frame_batch)
|
|
except NotImplementedError as e:
|
|
self.failure_count += 1
|
|
if self.failure_count >= FAIL_THRESHOLD:
|
|
LOG.error(
|
|
"NLF inference unsupported on device=%s after %d frames: %s. "
|
|
"TorchScript checkpoint is CUDA-only; install CUDA or switch backend.",
|
|
device, self.failure_count, e,
|
|
)
|
|
return
|
|
time.sleep(self.period)
|
|
continue
|
|
except Exception as e:
|
|
self.failure_count += 1
|
|
if self.failure_count >= FAIL_THRESHOLD:
|
|
LOG.error("NLF inference failed %d frames in a row, stopping: %s",
|
|
self.failure_count, e)
|
|
return
|
|
LOG.warning("inference failed: %s", e)
|
|
time.sleep(self.period)
|
|
continue
|
|
|
|
self._record_success()
|
|
|
|
verts_all = pred.get("vertices3d_nonparam")
|
|
joints_all = pred.get("joints3d_nonparam")
|
|
trans_all = pred.get("trans")
|
|
|
|
if verts_all is None or len(verts_all) == 0:
|
|
with self.state.lock():
|
|
self.state.persons_nlf = []
|
|
time.sleep(self.period)
|
|
continue
|
|
|
|
verts_batch = verts_all[0]
|
|
joints_batch = joints_all[0] if joints_all is not None else None
|
|
trans_batch = trans_all[0] if trans_all is not None else None
|
|
|
|
n_detected = min(verts_batch.shape[0], self.num_persons)
|
|
t_now = time.monotonic()
|
|
|
|
bboxes = []
|
|
for i in range(n_detected):
|
|
v = verts_batch[i].cpu().numpy()
|
|
xmin, ymin = v[:, 0].min(), v[:, 1].min()
|
|
xmax, ymax = v[:, 0].max(), v[:, 1].max()
|
|
bboxes.append([PoseKp(x=float(xmin), y=float(ymin), c=1.0),
|
|
PoseKp(x=float(xmax), y=float(ymax), c=1.0)])
|
|
|
|
ids = tracker.update(bboxes)
|
|
|
|
persons = []
|
|
for i in range(n_detected):
|
|
pid = ids[i] if i < len(ids) else i
|
|
if pid < 0:
|
|
continue
|
|
|
|
v_np = verts_batch[i].cpu().numpy()
|
|
j_np = (joints_batch[i].cpu().numpy()
|
|
if joints_batch is not None
|
|
else np.zeros((N_JOINTS, 3), dtype=np.float32))
|
|
t_np = (trans_batch[i].cpu().numpy()
|
|
if trans_batch is not None
|
|
else np.zeros(3, dtype=np.float32))
|
|
|
|
pid_c = pid % self.num_persons
|
|
t_smooth = np.array([
|
|
self._smooth_pos[pid_c][k](float(t_np[k]), t_now)
|
|
for k in range(3)
|
|
], dtype=np.float32)
|
|
|
|
persons.append(NLFPerson(
|
|
pid=int(pid),
|
|
vertices_3d=tuple(map(tuple, v_np)),
|
|
joints_3d=tuple(map(tuple, j_np)),
|
|
translation=tuple(t_smooth.tolist()),
|
|
confidence=1.0,
|
|
))
|
|
|
|
with self.state.lock():
|
|
self.state.persons_nlf = persons
|
|
self.state.nlf_last_t = t_now
|
|
|
|
dt = time.monotonic() - t0
|
|
if dt < self.period:
|
|
time.sleep(self.period - dt)
|
|
|
|
cap.release()
|
|
LOG.info("nlf worker stopped")
|