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:
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user