feat(detrpose): expose model size as CLI flag

DETRPose ships in three sizes (n/s/l); previously switching required
editing the DEFAULT_MODEL_SIZE constant in source. Add _VALID_SIZES and
_configure_model_size() to DETRPoseWorker for validated size selection,
and wire --detrpose-model-size into the argparse CLI in main.py.
This commit is contained in:
L'électron rare
2026-05-13 12:31:11 +02:00
parent a1b3ba1f91
commit 18368b2ded
3 changed files with 40 additions and 2 deletions
+10 -1
View File
@@ -72,6 +72,7 @@ CACHE_DIR = Path.home() / ".cache" / "av-live-detrpose"
REPO_DIR = CACHE_DIR / "DETRPose"
# Modele N (nano) par defaut : 16 MB, le plus rapide.
DEFAULT_MODEL_SIZE = "n"
_VALID_SIZES = {"n", "s", "l"}
DEFAULT_CKPT = CACHE_DIR / f"detrpose_hgnetv2_{DEFAULT_MODEL_SIZE}.pth"
DEFAULT_CONFIG_REL = f"configs/detrpose/detrpose_hgnetv2_{DEFAULT_MODEL_SIZE}.py"
@@ -124,7 +125,7 @@ class DETRPoseWorker:
self.period = 1.0 / max(1.0, target_fps)
self.num_persons = num_persons
self.score_thresh = score_thresh
self.model_size = model_size
self._configure_model_size(model_size)
self.device_pref = device
self._stop = threading.Event()
self._thread: threading.Thread | None = None
@@ -137,6 +138,14 @@ class DETRPoseWorker:
def stop(self) -> None:
self._stop.set()
def _configure_model_size(self, size: str) -> None:
"""Validate and set model_size; raise ValueError for unknown sizes."""
if size not in _VALID_SIZES:
raise ValueError(
f"DETRPose model_size must be one of {sorted(_VALID_SIZES)}, got {size!r}"
)
self.model_size = size
# ------------------------------------------------------------------
# Chargement modele : on importe le repo DETRPose en ajoutant son
# dossier au sys.path, puis on suit le pattern de tools/inference/torch_inf.py.
+6 -1
View File
@@ -297,7 +297,8 @@ class AppDelegate(NSObject):
from .detrpose import DETRPoseWorker, is_available
if is_available():
self._pose_worker = DETRPoseWorker(
self._state, num_persons=4)
self._state, num_persons=4,
model_size=getattr(self._opts, "detrpose_model_size", "n"))
self._pose_worker.start()
LOG.info("worker: DETRPose (transformer, body 17 kp × 4)")
return
@@ -495,6 +496,10 @@ def main() -> int:
p.add_argument("--pose-device", default="mps",
choices=("cpu", "mps", "cuda:0"),
help="Device YOLO inference (default mps)")
p.add_argument("--detrpose-model-size",
choices=["n", "s", "l"],
default="n",
help="DETRPose model size (default: n)")
p.add_argument("-v", "--verbose", action="store_true")
opts = p.parse_args()
+24
View File
@@ -0,0 +1,24 @@
"""DETRPose model size flows from CLI through to the worker."""
import pytest
from data_only_viz.detrpose import DETRPoseWorker, DEFAULT_MODEL_SIZE
def test_default_model_size_unchanged():
assert DEFAULT_MODEL_SIZE == "n"
def test_worker_accepts_model_size_kwarg():
# Should not raise; we don't load the model (that requires the extra)
worker = DETRPoseWorker.__new__(DETRPoseWorker)
worker.model_size = None
worker._configure_model_size("s")
assert worker.model_size == "s"
def test_worker_rejects_invalid_model_size():
worker = DETRPoseWorker.__new__(DETRPoseWorker)
worker.model_size = None
with pytest.raises(ValueError):
worker._configure_model_size("xxxl")