From e96c956a6f4f29497016c77c6a710694f8fa3549 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 13:50:27 +0200 Subject: [PATCH 01/23] Add pose backend interface and profiling flag Introduces a new `ENABLE_MODELS_PROFILING` config toggle and adds an abstract `PoseBackend` base interface in inference services. This establishes a shared contract (`init_inference`, `get_pose`, `make_pose_packet`, `close`) for backend implementations and prepares the codebase for consistent backend behavior and optional profiling instrumentation. --- dlclivegui/config.py | 1 + dlclivegui/services/inference/base.py | 34 +++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/dlclivegui/config.py b/dlclivegui/config.py index 14d00250..a48333b6 100644 --- a/dlclivegui/config.py +++ b/dlclivegui/config.py @@ -42,6 +42,7 @@ DEBUG_TRIGGER_LOGS = False ### Extra logs for DLC lifecycle (model loading, etc) DLC_LIFECYCLE_EXTRA_LOGS: bool = False +ENABLE_MODELS_PROFILING: bool = False # MAIN_WINDOW_DO_LOG_TIMING: bool = False #### Backends BASLER_DO_LOG_TIMING: bool = False diff --git a/dlclivegui/services/inference/base.py b/dlclivegui/services/inference/base.py index 17d0286b..918ef554 100644 --- a/dlclivegui/services/inference/base.py +++ b/dlclivegui/services/inference/base.py @@ -1,3 +1,5 @@ +import logging +from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum, auto from typing import Any @@ -6,6 +8,8 @@ from dlclivegui.config import ModelType +logger = logging.getLogger(__name__) + class PoseBackends(Enum): DLC_LIVE = auto() @@ -63,3 +67,33 @@ class ProcessorStats: # Separated timing for GPU vs socket processor avg_gpu_inference_time: float = 0.0 # Pure model inference avg_processor_overhead: float = 0.0 # Socket processor overhead + + +class PoseBackend(ABC): + """Common interface for pose-estimation backends.""" + + @abstractmethod + def init_inference( + self, + init_frame: np.ndarray, + ) -> None: + """Initialize inference using the first input frame.""" + + @abstractmethod + def get_pose( + self, + frame: np.ndarray, + frame_time: float | None = None, + ) -> np.ndarray | None: + """Return pose data for one frame.""" + + @abstractmethod + def make_pose_packet( + self, + pose: np.ndarray | None, + ) -> PosePacket: + """Wrap pose data and backend metadata for consumers.""" + + @abstractmethod + def close(self) -> None: + """Release backend resources.""" From fd225b8193d2f21a68ed630616215c0f8b341c43 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 13:50:48 +0200 Subject: [PATCH 02/23] Add threaded pose inference processor Introduce a new `PoseProcessor` service for background pose estimation with a pluggable backend. It runs inference in a dedicated worker thread, uses a bounded queue to handle frame backpressure, emits Qt signals for results/lifecycle/errors, and tracks runtime stats (latency, FPS, dropped frames, and optional profiling timings). The worker also supports graceful stop/reset/shutdown behavior and backend cleanup. --- dlclivegui/services/inference/processor.py | 270 +++++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 dlclivegui/services/inference/processor.py diff --git a/dlclivegui/services/inference/processor.py b/dlclivegui/services/inference/processor.py new file mode 100644 index 00000000..a4dfbb14 --- /dev/null +++ b/dlclivegui/services/inference/processor.py @@ -0,0 +1,270 @@ +import logging +import queue +import threading +import time +from collections import deque +from collections.abc import Callable +from typing import Any + +import numpy as np +from qtpy.QtCore import QObject, Signal + +from dlclivegui.config import ENABLE_MODELS_PROFILING + +from .base import PoseResult, ProcessorStats + +logger = logging.getLogger(__name__) + + +class PoseProcessor(QObject): + """ + Background pose estimation using a pluggable backend. + + backend_factory: () -> backend implementing init_inference() and get_pose() + """ + + pose_ready = Signal(object) + error = Signal(str) + initialized = Signal(bool) + frame_processed = Signal() + + def __init__(self) -> None: + super().__init__() + self._backend_factory: Callable[[], Any] | None = None + self._backend: Any | None = None + + self._queue: queue.Queue[Any] | None = None + self._worker_thread: threading.Thread | None = None + self._stop_event = threading.Event() + + self._frames_enqueued = 0 + self._frames_processed = 0 + self._frames_dropped = 0 + self._latencies: deque[float] = deque(maxlen=60) + self._processing_times: deque[float] = deque(maxlen=60) + self._queue_wait_times: deque[float] = deque(maxlen=60) + self._inference_times: deque[float] = deque(maxlen=60) + self._signal_emit_times: deque[float] = deque(maxlen=60) + self._total_process_times: deque[float] = deque(maxlen=60) + self._stats_lock = threading.Lock() + + def configure(self, backend_factory: Callable[[], Any]) -> None: + self._backend_factory = backend_factory + + def is_configured(self) -> bool: + return self._backend_factory is not None + + def reset(self) -> None: + self._stop_worker() + self._backend = None + with self._stats_lock: + self._frames_enqueued = 0 + self._frames_processed = 0 + self._frames_dropped = 0 + self._latencies.clear() + self._processing_times.clear() + self._queue_wait_times.clear() + self._inference_times.clear() + self._signal_emit_times.clear() + self._total_process_times.clear() + + def shutdown(self) -> None: + self.reset() + + def enqueue_frame(self, frame: np.ndarray, timestamp: float) -> None: + if self._worker_thread is None: + self._start_worker(frame.copy(), timestamp) + return + + if self._queue is None: + return + + try: + self._queue.put_nowait((frame.copy(), timestamp, time.perf_counter())) + with self._stats_lock: + self._frames_enqueued += 1 + except queue.Full: + with self._stats_lock: + self._frames_dropped += 1 + + def get_stats(self) -> ProcessorStats: + queue_size = self._queue.qsize() if self._queue is not None else 0 + with self._stats_lock: + avg_latency = sum(self._latencies) / len(self._latencies) if self._latencies else 0.0 + last_latency = self._latencies[-1] if self._latencies else 0.0 + + if len(self._processing_times) >= 2: + duration = self._processing_times[-1] - self._processing_times[0] + processing_fps = (len(self._processing_times) - 1) / duration if duration > 0 else 0.0 + else: + processing_fps = 0.0 + + avg_queue_wait = ( + sum(self._queue_wait_times) / len(self._queue_wait_times) if self._queue_wait_times else 0.0 + ) + avg_inference = sum(self._inference_times) / len(self._inference_times) if self._inference_times else 0.0 + avg_signal_emit = ( + sum(self._signal_emit_times) / len(self._signal_emit_times) if self._signal_emit_times else 0.0 + ) + avg_total = ( + sum(self._total_process_times) / len(self._total_process_times) if self._total_process_times else 0.0 + ) + + return ProcessorStats( + frames_enqueued=self._frames_enqueued, + frames_processed=self._frames_processed, + frames_dropped=self._frames_dropped, + queue_size=queue_size, + processing_fps=processing_fps, + average_latency=avg_latency, + last_latency=last_latency, + avg_queue_wait=avg_queue_wait, + avg_inference_time=avg_inference, + avg_signal_emit_time=avg_signal_emit, + avg_total_process_time=avg_total, + ) + + def _start_worker(self, init_frame: np.ndarray, init_timestamp: float) -> None: + if self._worker_thread is not None and self._worker_thread.is_alive(): + return + self._queue = queue.Queue(maxsize=1) + self._stop_event.clear() + self._worker_thread = threading.Thread( + target=self._worker_loop, + args=(init_frame, init_timestamp), + name="PoseWorker", + daemon=True, + ) + self._worker_thread.start() + + def _stop_worker(self) -> None: + worker = self._worker_thread + if worker is None: + return + + self._stop_event.set() + worker.join(timeout=2.0) + + if worker.is_alive(): + logger.warning("Pose worker did not stop within the timeout.") + return + + self._worker_thread = None + self._queue = None + + backend = self._backend + self._backend = None + + if backend is not None: + try: + backend.close() + except Exception: + logger.exception("Failed to close pose backend.") + + def _process_frame( + self, + frame: np.ndarray, + timestamp: float, + enqueue_time: float, + *, + queue_wait_time: float, + ) -> None: + backend = self._backend + if backend is None: + raise RuntimeError("Pose backend is not initialized.") + + inference_started = time.perf_counter() + + pose = backend.get_pose( + frame, + frame_time=timestamp, + ) + packet = backend.make_pose_packet(pose) + + inference_time = time.perf_counter() - inference_started + + signal_started = time.perf_counter() + self.pose_ready.emit( + PoseResult( + pose=pose, + timestamp=timestamp, + packet=packet, + ) + ) + signal_time = time.perf_counter() - signal_started + + finished = time.perf_counter() + latency = finished - enqueue_time + total_time = finished - enqueue_time + + with self._stats_lock: + self._frames_processed += 1 + self._latencies.append(latency) + self._processing_times.append(finished) + + if ENABLE_MODELS_PROFILING: + self._queue_wait_times.append(queue_wait_time) + self._inference_times.append(inference_time) + self._signal_emit_times.append(signal_time) + self._total_process_times.append(total_time) + + self.frame_processed.emit() + + def _worker_loop(self, init_frame: np.ndarray, init_timestamp: float) -> None: + try: + if self._backend_factory is None: + raise RuntimeError("No backend configured.") + self._backend = self._backend_factory() + + self._backend.init_inference(init_frame) + self.initialized.emit(True) + + self._process_frame(init_frame, init_timestamp, time.perf_counter(), queue_wait_time=0.0) + with self._stats_lock: + self._frames_enqueued += 1 + + except Exception as exc: + logger.exception("Failed to initialize pose backend", exc_info=exc) + self.error.emit(str(exc)) + self.initialized.emit(False) + return + + while True: + if self._stop_event.is_set(): + if self._queue is not None: + try: + frame, ts, enq = self._queue.get_nowait() + except queue.Empty: + break + else: + try: + self._process_frame(frame, ts, enq, queue_wait_time=0.0) + except Exception as exc: + logger.exception("Pose inference failed", exc_info=exc) + self.error.emit(str(exc)) + finally: + try: + self._queue.task_done() + except ValueError: + pass + continue + + try: + wait_start = time.perf_counter() + frame, ts, enq = self._queue.get(timeout=0.05) + qwait = time.perf_counter() - wait_start + except queue.Empty: + continue + + try: + self._process_frame(frame, ts, enq, queue_wait_time=qwait) + except Exception as exc: + logger.exception("Pose inference failed", exc_info=exc) + self.error.emit(str(exc)) + finally: + try: + self._queue.task_done() + except ValueError: + pass + + logger.info("Pose worker thread exiting") From 5c9364725f719fe73a86b9a49998b295be84c299 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 13:51:06 +0200 Subject: [PATCH 03/23] Add POET backend, skeleton, and weights Introduce a new POET inference module under `services/inference/models/poet`, including a `POETBackend` for model initialization and pose decoding, COCO-17 keypoint/skeleton definitions, and a Qt worker to download/cache default POET weights. This lays the groundwork for running POET-based multi-person pose estimation in the GUI. --- .../services/inference/models/__init__.py | 0 .../inference/models/poet/__init__.py | 0 .../inference/models/poet/poet_processor.py | 241 ++++++++++++++++++ .../inference/models/poet/skeleton.py | 53 ++++ .../services/inference/models/poet/weights.py | 63 +++++ 5 files changed, 357 insertions(+) create mode 100644 dlclivegui/services/inference/models/__init__.py create mode 100644 dlclivegui/services/inference/models/poet/__init__.py create mode 100644 dlclivegui/services/inference/models/poet/poet_processor.py create mode 100644 dlclivegui/services/inference/models/poet/skeleton.py create mode 100644 dlclivegui/services/inference/models/poet/weights.py diff --git a/dlclivegui/services/inference/models/__init__.py b/dlclivegui/services/inference/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dlclivegui/services/inference/models/poet/__init__.py b/dlclivegui/services/inference/models/poet/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dlclivegui/services/inference/models/poet/poet_processor.py b/dlclivegui/services/inference/models/poet/poet_processor.py new file mode 100644 index 00000000..44a3d78e --- /dev/null +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -0,0 +1,241 @@ +"""POET pose-estimation backend.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from dlclivegui.services.dlc_processor import ( + PoseBackends, + PosePacket, + PoseSource, +) +from dlclivegui.services.inference.base import PoseBackend + +from .skeleton import POET_KEYPOINT_NAMES, POET_SKELETON_EDGES + +try: + from poet_live import POET, PostProcess + from poet_live.models.backbone import Backbone, Joiner + from poet_live.models.position_encoding import PositionEmbeddingSine + from poet_live.models.transformer import Transformer +except ImportError as exc: + raise ImportError("POET imports failed. Ensure the POET package and its dependencies are installed.") from exc + + +logger = logging.getLogger(__name__) + + +class POETBackend(PoseBackend): + """POET pose-inference backend using COCO-17 keypoints.""" + + def __init__( + self, + checkpoint_path: str, + *, + device: str = "auto", + threshold: float = 0.7, + use_amp: bool = True, + ) -> None: + checkpoint = Path(checkpoint_path).expanduser() + + if not checkpoint.is_file(): + raise FileNotFoundError(f"POET checkpoint not found: {checkpoint}") + + if checkpoint.suffix.lower() not in {".pt", ".pth"}: + raise ValueError("POET checkpoint must use a .pt or .pth extension.") + + self._checkpoint_path = checkpoint + self._requested_device = device + self._threshold = float(threshold) + self._use_amp = bool(use_amp) + + self._model: Any | None = None + self._postprocessor: Any | None = None + self._device: torch.device | None = None + + self._mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + self._std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + + @staticmethod + def _resolve_torch_device( + device: str | None, + ) -> torch.device: + requested = device.strip().lower() if device else "auto" + + if requested in {"auto", "best"}: + if torch.cuda.is_available(): + return torch.device("cuda") + + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return torch.device("mps") + + return torch.device("cpu") + + if requested.startswith("cuda") and not torch.cuda.is_available(): + logger.warning( + "Requested device %r but CUDA is unavailable; using CPU.", + device, + ) + return torch.device("cpu") + + return torch.device(requested) + + def init_inference( + self, + init_frame: np.ndarray, + ) -> None: + ( + self._model, + self._postprocessor, + self._device, + ) = self._build_model() + + # Warm up without emitting a result. PoseProcessor emits the + # initial frame's result after init_inference() returns. + self.get_pose(init_frame) + + def _build_model( + self, + ) -> tuple[Any, Any, torch.device]: + device = self._resolve_torch_device(self._requested_device) + hidden_dimension = 256 + + backbone = Backbone( + "resnet50", + train_backbone=False, + return_interm_layers=False, + dilation5=False, + dilation4=False, + ) + position_encoding = PositionEmbeddingSine( + hidden_dimension // 2, + normalize=True, + ) + joined_backbone = Joiner( + backbone, + position_encoding, + ) + joined_backbone.num_channels = backbone.num_channels + + transformer = Transformer( + d_model=hidden_dimension, + return_intermediate_dec=True, + ) + model = POET( + joined_backbone, + transformer, + num_classes=2, + num_queries=25, + aux_loss=False, + ) + + checkpoint = torch.load( + self._checkpoint_path, + map_location="cpu", + weights_only=False, + ) + model.load_state_dict( + checkpoint["model"], + strict=True, + ) + + model.to(device).eval() + postprocessor = PostProcess().to(device) + + self._mean = self._mean.to(device) + self._std = self._std.to(device) + + return model, postprocessor, device + + @torch.no_grad() + def get_pose( + self, + frame: np.ndarray, + frame_time: float | None = None, + ) -> np.ndarray | None: + del frame_time + + model = self._model + postprocessor = self._postprocessor + device = self._device + + if model is None or postprocessor is None or device is None: + raise RuntimeError("POET backend is not initialized.") + + rgb = frame[..., ::-1].copy() + height, width = rgb.shape[:2] + + image = torch.from_numpy(rgb).to(device).permute(2, 0, 1).float().unsqueeze(0) / 255.0 + image = (image - self._mean) / self._std + + use_amp = self._use_amp and device.type == "cuda" + + with torch.autocast( + device_type=device.type, + enabled=use_amp, + ): + outputs = model(image) + + target_sizes = torch.tensor( + [[height, width]], + device=device, + ) + result = postprocessor( + outputs, + target_sizes=target_sizes, + )[0] + + scores = result["scores"] + keypoints = result["keypoints"] + + keep = scores >= self._threshold + if keep.sum().item() == 0: + return None + + scores = scores[keep] + keypoints = keypoints[keep].reshape( + -1, + len(POET_KEYPOINT_NAMES), + 3, + ) + + keypoints = keypoints.clone() + keypoints[:, :, 2] = scores[:, None].clamp(0, 1) + + ordering = torch.argsort( + scores, + descending=True, + ) + keypoints = keypoints[ordering] + + return keypoints.detach().cpu().numpy().astype(np.float32) + + def make_pose_packet( + self, + pose: np.ndarray | None, + ) -> PosePacket: + individual_count = pose.shape[0] if pose is not None and pose.ndim == 3 else 0 + + return PosePacket( + schema_version=1, + keypoints=pose, + keypoint_names=list(POET_KEYPOINT_NAMES), + individual_ids=[f"person_{index}" for index in range(individual_count)], + skeleton_id="poet.coco17", + skeleton_edges=POET_SKELETON_EDGES, + source=PoseSource( + backend=PoseBackends.POET, + model_type=None, + ), + raw=pose, + ) + + def close(self) -> None: + self._model = None + self._postprocessor = None + self._device = None diff --git a/dlclivegui/services/inference/models/poet/skeleton.py b/dlclivegui/services/inference/models/poet/skeleton.py new file mode 100644 index 00000000..06487d20 --- /dev/null +++ b/dlclivegui/services/inference/models/poet/skeleton.py @@ -0,0 +1,53 @@ +# dlclivegui/services/inference/models/poet/skeleton.py +from __future__ import annotations + +from dlclivegui.display.skeleton import ( + SkeletonDefinition, + SkeletonEdge, +) + +POET_KEYPOINT_NAMES = ( + "nose", + "left_eye", + "right_eye", + "left_ear", + "right_ear", + "left_shoulder", + "right_shoulder", + "left_elbow", + "right_elbow", + "left_wrist", + "right_wrist", + "left_hip", + "right_hip", + "left_knee", + "right_knee", + "left_ankle", + "right_ankle", +) + + +POET_SKELETON = SkeletonDefinition( + identifier="poet.coco17", + display_name="POET COCO-17", + edges=( + SkeletonEdge("left_ear", "left_eye"), + SkeletonEdge("right_ear", "right_eye"), + SkeletonEdge("left_eye", "nose"), + SkeletonEdge("nose", "right_eye"), + SkeletonEdge("left_shoulder", "right_shoulder"), + SkeletonEdge("left_shoulder", "left_elbow"), + SkeletonEdge("right_shoulder", "right_elbow"), + SkeletonEdge("nose", "left_shoulder"), + SkeletonEdge("nose", "right_shoulder"), + SkeletonEdge("left_shoulder", "left_hip"), + SkeletonEdge("right_shoulder", "right_hip"), + SkeletonEdge("left_elbow", "left_wrist"), + SkeletonEdge("right_elbow", "right_wrist"), + SkeletonEdge("left_hip", "right_hip"), + SkeletonEdge("left_hip", "left_knee"), + SkeletonEdge("right_hip", "right_knee"), + SkeletonEdge("left_knee", "left_ankle"), + SkeletonEdge("right_knee", "right_ankle"), + ), +) diff --git a/dlclivegui/services/inference/models/poet/weights.py b/dlclivegui/services/inference/models/poet/weights.py new file mode 100644 index 00000000..9ac1b53f --- /dev/null +++ b/dlclivegui/services/inference/models/poet/weights.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import os +import tempfile +import urllib.request +from pathlib import Path + +from PySide6.QtCore import QObject, Signal + + +def poet_default_weights_dir() -> Path: + return Path(tempfile.gettempdir()) / "dlclivegui" / "poet" + + +POET_WEIGHTS_URL = "https://zenodo.org/records/7972042/files/poet_ckpt.pth?download=1" +POET_WEIGHTS_FILENAME = "poet_resnet50.pth" + + +class WeightsDownloadWorker(QObject): + progress = Signal(int) # 0..100 + finished = Signal(str) # path + error = Signal(str) + + def __init__(self, url: str, dest: Path): + super().__init__() + self.url = url + self.dest = dest + + def run(self) -> None: + try: + self.dest.parent.mkdir(parents=True, exist_ok=True) + if self.dest.is_file(): + self.progress.emit(100) + self.finished.emit(str(self.dest)) + return + tmp = self.dest.with_suffix(self.dest.suffix + ".part") + + req = urllib.request.Request(self.url, headers={"User-Agent": "DLCLiveGUI"}) + with urllib.request.urlopen(req) as resp, open(tmp, "wb") as f: + total = resp.length or 0 + done = 0 + chunk = 1024 * 256 + + while True: + buf = resp.read(chunk) + if not buf: + break + f.write(buf) + done += len(buf) + if total > 0: + self.progress.emit(int(done * 100 / total)) + + final_path = self.dest.parent / POET_WEIGHTS_FILENAME + os.replace(tmp, final_path) + self.progress.emit(100) + self.finished.emit(str(final_path)) + + except Exception as e: + try: + tmp.unlink(missing_ok=True) + except Exception: + pass + self.error.emit(str(e)) From 6b54d22573f458fb6cacda840ac885ef8a848629 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 14:03:36 +0200 Subject: [PATCH 04/23] Harden POET backend and inference lifecycle Refactors inference plumbing to better support POET and improve runtime safety. This adds a string-based `PoseBackends` enum with `POET`, renames the model profiling config flag, tightens `PoseProcessor` typing/Qt imports, and makes worker shutdown/reset and backend-init failure cleanup explicit. It also lazy-loads POET dependencies with a clearer runtime error, decouples POET skeleton metadata from display-layer types, and fixes weight download finalization to atomically replace the configured destination while logging partial-file cleanup failures. --- dlclivegui/config.py | 2 +- dlclivegui/services/inference/base.py | 5 +- .../inference/models/poet/poet_processor.py | 26 +++++---- .../inference/models/poet/skeleton.py | 53 +++++++++---------- .../services/inference/models/poet/weights.py | 24 +++++---- dlclivegui/services/inference/processor.py | 36 +++++++++---- 6 files changed, 82 insertions(+), 64 deletions(-) diff --git a/dlclivegui/config.py b/dlclivegui/config.py index a48333b6..0558e7cc 100644 --- a/dlclivegui/config.py +++ b/dlclivegui/config.py @@ -42,7 +42,7 @@ DEBUG_TRIGGER_LOGS = False ### Extra logs for DLC lifecycle (model loading, etc) DLC_LIFECYCLE_EXTRA_LOGS: bool = False -ENABLE_MODELS_PROFILING: bool = False +MODEL_INFERENCE_PROFILING_ENABLED: bool = False # MAIN_WINDOW_DO_LOG_TIMING: bool = False #### Backends BASLER_DO_LOG_TIMING: bool = False diff --git a/dlclivegui/services/inference/base.py b/dlclivegui/services/inference/base.py index 918ef554..891e5792 100644 --- a/dlclivegui/services/inference/base.py +++ b/dlclivegui/services/inference/base.py @@ -11,8 +11,9 @@ logger = logging.getLogger(__name__) -class PoseBackends(Enum): - DLC_LIVE = auto() +class PoseBackends(str, Enum): + DLC_LIVE = "DLC_LIVE" + POET = "POET" class WorkerState(Enum): diff --git a/dlclivegui/services/inference/models/poet/poet_processor.py b/dlclivegui/services/inference/models/poet/poet_processor.py index 44a3d78e..4afa2a5c 100644 --- a/dlclivegui/services/inference/models/poet/poet_processor.py +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -9,24 +9,15 @@ import numpy as np import torch -from dlclivegui.services.dlc_processor import ( +from dlclivegui.services.inference.base import ( + PoseBackend, PoseBackends, PosePacket, PoseSource, ) -from dlclivegui.services.inference.base import PoseBackend from .skeleton import POET_KEYPOINT_NAMES, POET_SKELETON_EDGES -try: - from poet_live import POET, PostProcess - from poet_live.models.backbone import Backbone, Joiner - from poet_live.models.position_encoding import PositionEmbeddingSine - from poet_live.models.transformer import Transformer -except ImportError as exc: - raise ImportError("POET imports failed. Ensure the POET package and its dependencies are installed.") from exc - - logger = logging.getLogger(__name__) @@ -102,7 +93,20 @@ def init_inference( def _build_model( self, ) -> tuple[Any, Any, torch.device]: + try: + from poet_live import POET, PostProcess + from poet_live.models.backbone import Backbone, Joiner + from poet_live.models.position_encoding import ( + PositionEmbeddingSine, + ) + from poet_live.models.transformer import Transformer + except ImportError as exc: + raise RuntimeError( + "POET is not installed. Install the POET dependencies before using this pose backend." + ) from exc + device = self._resolve_torch_device(self._requested_device) + hidden_dimension = 256 backbone = Backbone( diff --git a/dlclivegui/services/inference/models/poet/skeleton.py b/dlclivegui/services/inference/models/poet/skeleton.py index 06487d20..e3284589 100644 --- a/dlclivegui/services/inference/models/poet/skeleton.py +++ b/dlclivegui/services/inference/models/poet/skeleton.py @@ -1,10 +1,6 @@ -# dlclivegui/services/inference/models/poet/skeleton.py -from __future__ import annotations +"""POET COCO-17 keypoint and skeleton metadata.""" -from dlclivegui.display.skeleton import ( - SkeletonDefinition, - SkeletonEdge, -) +from __future__ import annotations POET_KEYPOINT_NAMES = ( "nose", @@ -27,27 +23,26 @@ ) -POET_SKELETON = SkeletonDefinition( - identifier="poet.coco17", - display_name="POET COCO-17", - edges=( - SkeletonEdge("left_ear", "left_eye"), - SkeletonEdge("right_ear", "right_eye"), - SkeletonEdge("left_eye", "nose"), - SkeletonEdge("nose", "right_eye"), - SkeletonEdge("left_shoulder", "right_shoulder"), - SkeletonEdge("left_shoulder", "left_elbow"), - SkeletonEdge("right_shoulder", "right_elbow"), - SkeletonEdge("nose", "left_shoulder"), - SkeletonEdge("nose", "right_shoulder"), - SkeletonEdge("left_shoulder", "left_hip"), - SkeletonEdge("right_shoulder", "right_hip"), - SkeletonEdge("left_elbow", "left_wrist"), - SkeletonEdge("right_elbow", "right_wrist"), - SkeletonEdge("left_hip", "right_hip"), - SkeletonEdge("left_hip", "left_knee"), - SkeletonEdge("right_hip", "right_knee"), - SkeletonEdge("left_knee", "left_ankle"), - SkeletonEdge("right_knee", "right_ankle"), - ), +POET_SKELETON_ID = "poet.coco17" + + +POET_SKELETON_EDGES = ( + ("left_ear", "left_eye"), + ("right_ear", "right_eye"), + ("left_eye", "nose"), + ("nose", "right_eye"), + ("left_shoulder", "right_shoulder"), + ("left_shoulder", "left_elbow"), + ("right_shoulder", "right_elbow"), + ("nose", "left_shoulder"), + ("nose", "right_shoulder"), + ("left_shoulder", "left_hip"), + ("right_shoulder", "right_hip"), + ("left_elbow", "left_wrist"), + ("right_elbow", "right_wrist"), + ("left_hip", "right_hip"), + ("left_hip", "left_knee"), + ("right_hip", "right_knee"), + ("left_knee", "left_ankle"), + ("right_knee", "right_ankle"), ) diff --git a/dlclivegui/services/inference/models/poet/weights.py b/dlclivegui/services/inference/models/poet/weights.py index 9ac1b53f..d361ba83 100644 --- a/dlclivegui/services/inference/models/poet/weights.py +++ b/dlclivegui/services/inference/models/poet/weights.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging import os import tempfile import urllib.request @@ -7,15 +8,18 @@ from PySide6.QtCore import QObject, Signal - -def poet_default_weights_dir() -> Path: - return Path(tempfile.gettempdir()) / "dlclivegui" / "poet" +logger = logging.getLogger(__name__) POET_WEIGHTS_URL = "https://zenodo.org/records/7972042/files/poet_ckpt.pth?download=1" POET_WEIGHTS_FILENAME = "poet_resnet50.pth" +def poet_default_weights_dir() -> Path: + # TODO: move to an app cache rather than a temporary directory + return Path(tempfile.gettempdir()) / "dlclivegui" / "poet" + + class WeightsDownloadWorker(QObject): progress = Signal(int) # 0..100 finished = Signal(str) # path @@ -50,14 +54,14 @@ def run(self) -> None: if total > 0: self.progress.emit(int(done * 100 / total)) - final_path = self.dest.parent / POET_WEIGHTS_FILENAME - os.replace(tmp, final_path) + os.replace(tmp, self.dest) self.progress.emit(100) - self.finished.emit(str(final_path)) + self.finished.emit(str(self.dest)) except Exception as e: - try: - tmp.unlink(missing_ok=True) - except Exception: - pass + if tmp is not None: + try: + tmp.unlink(missing_ok=True) + except OSError: + logger.exception("Failed to remove partial POET weights.") self.error.emit(str(e)) diff --git a/dlclivegui/services/inference/processor.py b/dlclivegui/services/inference/processor.py index a4dfbb14..13155162 100644 --- a/dlclivegui/services/inference/processor.py +++ b/dlclivegui/services/inference/processor.py @@ -7,11 +7,11 @@ from typing import Any import numpy as np -from qtpy.QtCore import QObject, Signal +from PySide6.QtCore import QObject, Signal -from dlclivegui.config import ENABLE_MODELS_PROFILING +from dlclivegui.config import MODEL_INFERENCE_PROFILING_ENABLED -from .base import PoseResult, ProcessorStats +from .base import PoseBackend, PoseResult, ProcessorStats logger = logging.getLogger(__name__) @@ -30,8 +30,8 @@ class PoseProcessor(QObject): def __init__(self) -> None: super().__init__() - self._backend_factory: Callable[[], Any] | None = None - self._backend: Any | None = None + self._backend_factory: Callable[[], PoseBackend] | None = None + self._backend: PoseBackend | None = None self._queue: queue.Queue[Any] | None = None self._worker_thread: threading.Thread | None = None @@ -48,14 +48,15 @@ def __init__(self) -> None: self._total_process_times: deque[float] = deque(maxlen=60) self._stats_lock = threading.Lock() - def configure(self, backend_factory: Callable[[], Any]) -> None: + def configure(self, backend_factory: Callable[[], PoseBackend]) -> None: self._backend_factory = backend_factory def is_configured(self) -> bool: return self._backend_factory is not None def reset(self) -> None: - self._stop_worker() + if not self._stop_worker(): + raise RuntimeError("Failed to stop worker thread") self._backend = None with self._stats_lock: self._frames_enqueued = 0 @@ -137,17 +138,17 @@ def _start_worker(self, init_frame: np.ndarray, init_timestamp: float) -> None: ) self._worker_thread.start() - def _stop_worker(self) -> None: + def _stop_worker(self) -> bool: worker = self._worker_thread if worker is None: - return + return True self._stop_event.set() worker.join(timeout=2.0) if worker.is_alive(): logger.warning("Pose worker did not stop within the timeout.") - return + return False self._worker_thread = None self._queue = None @@ -160,6 +161,7 @@ def _stop_worker(self) -> None: backend.close() except Exception: logger.exception("Failed to close pose backend.") + return True def _process_frame( self, @@ -202,7 +204,7 @@ def _process_frame( self._latencies.append(latency) self._processing_times.append(finished) - if ENABLE_MODELS_PROFILING: + if MODEL_INFERENCE_PROFILING_ENABLED: self._queue_wait_times.append(queue_wait_time) self._inference_times.append(inference_time) self._signal_emit_times.append(signal_time) @@ -227,6 +229,18 @@ def _worker_loop(self, init_frame: np.ndarray, init_timestamp: float) -> None: logger.exception("Failed to initialize pose backend", exc_info=exc) self.error.emit(str(exc)) self.initialized.emit(False) + + backend = self._backend + self._backend = None + + if backend is not None: + try: + backend.close() + except Exception as exc: + logger.exception("Failed to close backend", exc_info=exc) + + self._queue = None + self._worker_thread = None return while True: From 02ef7fd393621896f8423fb945034f73fd6b2dcc Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 14:36:08 +0200 Subject: [PATCH 05/23] Add backend-specific pose processor factory Introduce `create_pose_processor()` in the inference processor module to return the correct processor implementation for each backend (`DLCLiveProcessor` for `dlc`, `PoseProcessor` for `poet`) with explicit backend typing and unsupported-backend validation. Also update POET result construction to use the shared `POET_SKELETON_ID` constant instead of an inline string, keeping skeleton identifiers consistent. --- .../inference/models/poet/poet_processor.py | 4 ++-- dlclivegui/services/inference/processor.py | 17 ++++++++++++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/dlclivegui/services/inference/models/poet/poet_processor.py b/dlclivegui/services/inference/models/poet/poet_processor.py index 4afa2a5c..7cd56b00 100644 --- a/dlclivegui/services/inference/models/poet/poet_processor.py +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -16,7 +16,7 @@ PoseSource, ) -from .skeleton import POET_KEYPOINT_NAMES, POET_SKELETON_EDGES +from .skeleton import POET_KEYPOINT_NAMES, POET_SKELETON_EDGES, POET_SKELETON_ID logger = logging.getLogger(__name__) @@ -230,7 +230,7 @@ def make_pose_packet( keypoints=pose, keypoint_names=list(POET_KEYPOINT_NAMES), individual_ids=[f"person_{index}" for index in range(individual_count)], - skeleton_id="poet.coco17", + skeleton_id=POET_SKELETON_ID, skeleton_edges=POET_SKELETON_EDGES, source=PoseSource( backend=PoseBackends.POET, diff --git a/dlclivegui/services/inference/processor.py b/dlclivegui/services/inference/processor.py index 13155162..93b1dc98 100644 --- a/dlclivegui/services/inference/processor.py +++ b/dlclivegui/services/inference/processor.py @@ -4,16 +4,18 @@ import time from collections import deque from collections.abc import Callable -from typing import Any +from typing import Any, Literal import numpy as np from PySide6.QtCore import QObject, Signal from dlclivegui.config import MODEL_INFERENCE_PROFILING_ENABLED +from dlclivegui.services.dlc_processor import DLCLiveProcessor from .base import PoseBackend, PoseResult, ProcessorStats logger = logging.getLogger(__name__) +PoseBackendName = Literal["dlc", "poet"] class PoseProcessor(QObject): @@ -282,3 +284,16 @@ def _worker_loop(self, init_frame: np.ndarray, init_timestamp: float) -> None: pass logger.info("Pose worker thread exiting") + + +def create_pose_processor( + backend: PoseBackendName, +) -> DLCLiveProcessor | PoseProcessor: + """Create the processor service used for a pose backend.""" + if backend == "dlc": + return DLCLiveProcessor() + + if backend == "poet": + return PoseProcessor() + + raise ValueError(f"Unsupported pose backend: {backend!r}.") From 48890244ad39e4130c6fe4dbae7df2d24c097f84 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 14:36:21 +0200 Subject: [PATCH 06/23] Add selectable pose backends in main window Refactors inference handling to use a backend-agnostic pose processor interface and adds a Models > Pose backend menu to switch between DLCLive and POET. Updates model browsing/validation, signal wiring, control enablement, stats text, and inference lifecycle logic so backend-specific behavior is applied consistently (including recording hooks and shutdown). Also adds POET configuration support and persists the selected backend in settings. --- dlclivegui/gui/main_window.py | 461 +++++++++++++++++++++++++++------- 1 file changed, 374 insertions(+), 87 deletions(-) diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index 6ecb496c..0a2325c9 100644 --- a/dlclivegui/gui/main_window.py +++ b/dlclivegui/gui/main_window.py @@ -9,6 +9,7 @@ import threading import time from pathlib import Path +from typing import TYPE_CHECKING import numpy as np from PySide6.QtCore import QRect, QSettings, Qt, QTimer, QUrl, Signal @@ -83,6 +84,7 @@ scan_processor_package, ) from ..services.dlc_processor import DLCLiveProcessor, PoseResult +from ..services.inference.processor import PoseBackendName, PoseProcessor, create_pose_processor from ..services.multi_camera_controller import MultiCameraController, MultiFrameData, get_camera_id, get_display_id from ..services.recording_manager import RecordingManager from ..utils.settings_store import DLCLiveGUISettingsStore, ModelPathStore @@ -97,6 +99,8 @@ from .theme import LOGO, LOGO_ALPHA, AppStyle, apply_theme logger = logging.getLogger("DLCLiveGUI") +if TYPE_CHECKING: + from ..services.inference.models.poet.poet_processor import POETBackend class DLCLiveMainWindow(QMainWindow): @@ -148,7 +152,11 @@ def __init__(self, config: ApplicationSettings | None = None): self._fps_tracker = FPSTracker() self._rec_manager = RecordingManager() - self._dlc = DLCLiveProcessor() + self._dlc = create_pose_processor("dlc") + self._pose_processors = { + "dlc": self._dlc, + } + self._backend_name: PoseBackendName = "dlc" self.multi_camera_controller = MultiCameraController() ### Time debug self._dlc_timing = WorkerTimingStats( @@ -372,7 +380,7 @@ def _setup_ui(self) -> None: self.setStatusBar(QStatusBar()) self._build_menus() - QTimer.singleShot(0, self._show_logo_and_text) + self.show_logo_timer = QTimer.singleShot(0, self._show_logo_and_text) def _build_stats_layout(self, stats_widget: QWidget) -> QGridLayout: stats_layout = QGridLayout(stats_widget) @@ -394,11 +402,11 @@ def _build_stats_layout(self, stats_widget: QWidget) -> QGridLayout: row += 1 # DLC - title_dlc = QLabel("DLC Processor:") + title_dlc = QLabel("Pose processor:") title_dlc.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Preferred) stats_layout.addWidget(title_dlc, row, 0, alignment=Qt.AlignTop) - self.dlc_stats_label = QLabel("DLC processor idle") + self.dlc_stats_label = QLabel("Pose processor idle") self.dlc_stats_label.setWordWrap(True) self.dlc_stats_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) stats_layout.addWidget(self.dlc_stats_label, row, 1, alignment=Qt.AlignTop) @@ -470,6 +478,34 @@ def _build_menus(self) -> None: self._apply_theme(self._current_style) self._init_theme_actions() + # Models menu and pose backend selection + models_menu = self.menuBar().addMenu("&Models") + backend_menu = models_menu.addMenu("Pose backend") + + self.action_backend_dlc = QAction( + "DLCLive", + self, + checkable=True, + ) + self.action_backend_poet = QAction( + "POET", + self, + checkable=True, + ) + + self._backend_action_group = QActionGroup(self) + self._backend_action_group.setExclusive(True) + self._backend_action_group.addAction(self.action_backend_dlc) + self._backend_action_group.addAction(self.action_backend_poet) + + backend_menu.addAction(self.action_backend_dlc) + backend_menu.addAction(self.action_backend_poet) + + self.action_backend_dlc.triggered.connect(lambda checked: checked and self._set_backend("dlc")) + self.action_backend_poet.triggered.connect(lambda checked: checked and self._set_backend("poet")) + + self.action_backend_dlc.setChecked(True) + # Help menu help_menu = self.menuBar().addMenu("&Help") @@ -497,8 +533,8 @@ def _build_camera_group(self) -> QGroupBox: return group def _build_dlc_group(self) -> QGroupBox: - group = QGroupBox("DLCLive") - form = QFormLayout(group) + self.inference_group = QGroupBox("DLCLive") + form = QFormLayout(self.inference_group) path_layout = QHBoxLayout() self.model_path_edit = QLineEdit() @@ -596,7 +632,7 @@ def _build_dlc_group(self) -> QGroupBox: # self.show_predictions_checkbox.setChecked(True) # form.addRow(self.show_predictions_checkbox) - return group + return self.inference_group def _build_recording_group(self) -> QGroupBox: """Build recording controls group.""" @@ -880,6 +916,46 @@ def _build_viz_group(self) -> QGroupBox: return group # ------------------------------------------------------------------ signals + def _connect_pose_processor_signals(self) -> None: + processor = self._active_pose_processor + + processor.pose_ready.connect( + self._on_pose_ready, + Qt.ConnectionType.UniqueConnection, + ) + processor.error.connect( + self._on_pose_error, + Qt.ConnectionType.UniqueConnection, + ) + processor.initialized.connect( + self._on_pose_processor_initialised, + Qt.ConnectionType.UniqueConnection, + ) + + def _disconnect_pose_processor_signals(self) -> None: + processor = self._active_pose_processor + + connections = ( + ( + processor.pose_ready, + self._on_pose_ready, + ), + ( + processor.error, + self._on_pose_error, + ), + ( + processor.initialized, + self._on_pose_processor_initialised, + ), + ) + + for signal, slot in connections: + try: + signal.disconnect(slot) + except RuntimeError: + pass + def _connect_signals(self) -> None: self.preview_button.clicked.connect(self._start_preview) self.stop_preview_button.clicked.connect(self._stop_preview) @@ -915,9 +991,7 @@ def _connect_signals(self) -> None: self.multi_camera_controller.initialization_failed.connect(self._on_multi_camera_initialization_failed) # DLC processor signals - self._dlc.pose_ready.connect(self._on_pose_ready) - self._dlc.error.connect(self._on_dlc_error) - self._dlc.initialized.connect(self._on_dlc_initialised) + self._connect_pose_processor_signals() self.dlc_camera_combo.currentIndexChanged.connect(self._on_dlc_camera_changed) self.dlc_camera_combo.currentTextChanged.connect(self.dlc_camera_combo.update_shrink_width) self.processor_combo.currentIndexChanged.connect(self._on_processor_selection_changed) @@ -1050,6 +1124,16 @@ def _apply_config(self, config: ApplicationSettings, *, restore_local_prefs: boo # Update recording path preview self._update_recording_path_preview() + def _current_dlc_settings( + self, + *, + allow_empty_model_path: bool, + ) -> DLCProcessorSettings: + if self._backend_name == "dlc": + return self._dlc_settings_from_ui(allow_empty_model_path=allow_empty_model_path) + + return self._config.dlc.model_copy(deep=True) + def _current_config(self, *, allow_empty_model_path=False) -> ApplicationSettings: multi_camera = self._config.multi_camera active_cameras = multi_camera.get_active_cameras() @@ -1066,7 +1150,7 @@ def _current_config(self, *, allow_empty_model_path=False) -> ApplicationSetting return ApplicationSettings( camera=camera, multi_camera=multi_camera, - dlc=self._dlc_settings_from_ui(allow_empty_model_path=allow_empty_model_path), + dlc=self._current_dlc_settings(allow_empty_model_path=allow_empty_model_path), recording=self._recording_settings_from_ui(), bbox=self._bbox_settings_from_ui(), visualization=self._visualization_settings_from_ui(), @@ -1293,59 +1377,69 @@ def _save_config_to_path(self, path: Path) -> bool: return True def _action_browse_model(self) -> None: - # Prefer persisted last-used directory, then config.dlc.model_directory, then home start_dir = self._model_path_store.suggest_start_dir(self._config.dlc.model_directory) preselect = self._model_path_store.suggest_selected_file() - dlg = QFileDialog(self, "Select DLCLive model file") - dlg.setFileMode(QFileDialog.FileMode.ExistingFile) - dlg.setNameFilters( - [ - "Model files (*.pt *.pth)", - "PyTorch models (*.pt *.pth)", - "TensorFlow models (*.pb)", - ] + dialog = QFileDialog( + self, + ("Select DeepLabCut model" if self._backend_name == "dlc" else "Select POET weights"), ) - dlg.setDirectory(start_dir) + dialog.setFileMode(QFileDialog.FileMode.ExistingFile) + + if self._backend_name == "poet": + dialog.setNameFilter("POET weights (*.pt *.pth)") + else: + dialog.setNameFilters( + [ + "Model files (*.pt *.pth *.pb)", + "PyTorch models (*.pt *.pth)", + "TensorFlow models (*.pb)", + ] + ) + + dialog.setDirectory(start_dir) - # Preselect last used model if it exists (optional but nice) if preselect: - dlg.selectFile(preselect) + dialog.selectFile(preselect) - if dlg.exec(): - selected = dlg.selectedFiles() - if not selected: - return - file_path = Path(selected[0]).expanduser() - if not file_path.exists(): - QMessageBox.warning(self, "File not found", f"The selected file does not exist:\n{file_path}") - return + if not dialog.exec(): + return - try: - if file_path.suffix == ".pb": - # For TensorFlow, DLCLive expects a directory, so we pass the parent directory for validation - model_check_path = file_path.parent - else: - model_check_path = file_path - # Raise if the model backend cannot be determined (invalid file or unsupported extension) - DLCLiveProcessor.get_model_backend(str(model_check_path)) - except FileNotFoundError as e: - QMessageBox.warning(self, "Model selection error", str(e)) - return - except ValueError as e: - QMessageBox.warning(self, "Model selection error", str(e)) - return - file_path = str(file_path) - self.model_path_edit.setText(file_path) + selected = dialog.selectedFiles() + if not selected: + return - # Persist model path + directory - self._model_path_store.save_if_valid(file_path) + file_path = Path(selected[0]).expanduser() - # Optional: update config so next startup uses this directory too - try: - self._config.dlc.model_directory = str(Path(file_path).parent) - except Exception: - pass + if not file_path.is_file(): + QMessageBox.warning( + self, + "File not found", + f"The selected file does not exist:\n{file_path}", + ) + return + + try: + if self._backend_name == "poet": + if file_path.suffix.lower() not in { + ".pt", + ".pth", + }: + raise ValueError("POET weights must use a .pt or .pth extension.") + else: + model_path = file_path.parent if file_path.suffix.lower() == ".pb" else file_path + DLCLiveProcessor.get_model_backend(str(model_path)) + except (FileNotFoundError, ValueError) as exc: + QMessageBox.warning( + self, + "Model selection error", + str(exc), + ) + return + + selected_path = str(file_path) + self.model_path_edit.setText(selected_path) + self._model_path_store.save_if_valid(selected_path) def _action_browse_directory(self) -> None: directory = QFileDialog.getExistingDirectory(self, "Select output directory", str(Path.home())) @@ -1410,7 +1504,8 @@ def _action_view_documentation(self) -> None: def _custom_processor_enabled(self) -> bool: return bool( - getattr(self, "use_custom_proc_checkbox", None) + self._backend_name == "dlc" + and getattr(self, "use_custom_proc_checkbox", None) and self.use_custom_proc_checkbox.isChecked() and self.processor_combo.currentData() is not None ) @@ -1882,7 +1977,7 @@ def _on_multi_frame_processing_ready(self, frame_data: MultiFrameData) -> None: frame = frame_data.frames[dlc_cam_id] timestamp = frame_data.timestamps.get(dlc_cam_id, time.time()) with self._dlc_timing.measure("enqueue_frame"): - self._dlc.enqueue_frame(frame, timestamp) + self._active_pose_processor.enqueue_frame(frame, timestamp) self._dlc_timing.note_frame() self._dlc_timing.maybe_log() @@ -1970,9 +2065,14 @@ def _start_multi_camera_recording(self) -> None: self._show_error("Failed to start recording.") return - self._processor_recording_context = self._build_processor_recording_context(run_dir) - self._processor_recording_started_notified = False - self._notify_processor_recording_started(self._processor_recording_context) + if self._backend_name == "dlc": + self._processor_recording_context = self._build_processor_recording_context(run_dir) + self._processor_recording_started_notified = False + self._notify_processor_recording_started(self._processor_recording_context) + else: + self._processor_recording_context = None + self._processor_recording_started_notified = False + self.multi_camera_controller.set_recording_sink(self._rec_manager.write_frame) self.multi_camera_controller.set_recording_frame_is_enabled(True) @@ -2025,7 +2125,102 @@ def worker(): daemon=True, ).start() - def _get_dlc_processor_instance(self): + @property + def _active_pose_processor( + self, + ): + return self._get_pose_processor(self._backend_name) + + def _get_pose_processor( + self, + backend: PoseBackendName, + ): + processor = self._pose_processors.get(backend) + + if processor is None: + processor = create_pose_processor(backend) + self._pose_processors[backend] = processor + + return processor + + def _set_backend( + self, + name: str, + ) -> None: + if name not in {"dlc", "poet"}: + raise ValueError(f"Unsupported pose backend: {name!r}.") + + backend: PoseBackendName = name + + if backend == self._backend_name: + self._sync_backend_actions() + self._update_backend_ui() + return + + if self._dlc_active: + self._show_warning("Stop pose inference before switching backends.") + self._sync_backend_actions() + return + + self._disconnect_pose_processor_signals() + + self._backend_name = backend + + # Construct the processor lazily before connecting its signals. + self._get_pose_processor(backend) + + self._last_pose = None + self._overlay_renderer.clear_runtime_state() + + self._connect_pose_processor_signals() + self._sync_backend_actions() + self._update_backend_ui() + + self.settings.setValue( + "app/backend", + backend, + ) + + self.statusBar().showMessage( + f"Pose backend set to {self._backend_display_name()}", + 3000, + ) + + def _backend_display_name(self) -> str: + if self._backend_name == "poet": + return "POET" + + return "DLCLive" + + def _sync_backend_actions(self) -> None: + self.action_backend_dlc.blockSignals(True) + self.action_backend_poet.blockSignals(True) + + try: + self.action_backend_dlc.setChecked(self._backend_name == "dlc") + self.action_backend_poet.setChecked(self._backend_name == "poet") + finally: + self.action_backend_dlc.blockSignals(False) + self.action_backend_poet.blockSignals(False) + + def _update_backend_ui(self) -> None: + is_dlc = self._backend_name == "dlc" + + self.model_path_edit.setPlaceholderText("/path/to/exported/model" if is_dlc else "/path/to/poet_weights.pth") + + for widget in ( + self.processor_folder_edit, + self.browse_processor_folder_button, + self.refresh_processors_button, + self.processor_combo, + self.use_custom_proc_checkbox, + ): + widget.setVisible(is_dlc) + + self.processor_toggle_row.setVisible(is_dlc and self.processor_combo.currentData() is not None) + self.inference_group.setTitle(self._backend_display_name()) + + def _get_dlc_custom_processor_instance(self): """Return the active custom DLC processor instance, if available.""" processor = getattr(self._dlc, "_processor", None) @@ -2050,7 +2245,7 @@ def _save_processor_data_if_available(self) -> None: Return values are only logged; failure should not crash the GUI. """ - processor = self._get_dlc_processor_instance() + processor = self._get_dlc_custom_processor_instance() if processor is None: logger.debug("Processor save skipped: no processor instance available.") @@ -2068,7 +2263,7 @@ def _save_processor_data_if_available(self) -> None: logger.exception("Processor save() failed.") def _notify_processor_recording_started(self, context: dict) -> bool: - processor = self._get_dlc_processor_instance() + processor = self._get_dlc_custom_processor_instance() if processor is None: return False @@ -2143,7 +2338,7 @@ def _final_processor_recording_context(self) -> dict: return context def _notify_processor_recording_stopped(self) -> None: - processor = self._get_dlc_processor_instance() + processor = self._get_dlc_custom_processor_instance() if processor is None: return @@ -2380,6 +2575,57 @@ def _configure_dlc(self) -> bool: self._model_path_store.save_if_valid(settings.model_path) return True + def _configure_poet(self) -> bool: + checkpoint_text = self.model_path_edit.text().strip() + checkpoint = Path(checkpoint_text).expanduser() + + if not checkpoint_text: + self._show_error("Please select POET weights before starting inference.") + return False + + if not checkpoint.is_file(): + self._show_error(f"POET weights were not found:\n{checkpoint}") + return False + + if checkpoint.suffix.lower() not in { + ".pt", + ".pth", + }: + self._show_error("POET weights must use a .pt or .pth extension.") + return False + + device = self._config.dlc.device or "auto" + threshold = self._p_cutoff + + def backend_factory() -> POETBackend: + + from ..services.inference.models.poet.poet_processor import POETBackend + + return POETBackend( + str(checkpoint), + device=device, + threshold=threshold, + use_amp=True, + ) + + processor = self._get_pose_processor("poet") + + if not isinstance(processor, PoseProcessor): + raise RuntimeError("The POET processor does not support configuration.") + + processor.configure(backend_factory) + return True + + def _configure_pose_backend(self) -> bool: + if self._backend_name == "dlc": + return self._configure_dlc() + + if self._backend_name == "poet": + return self._configure_poet() + + self._show_error(f"Unsupported pose backend: {self._backend_name!r}.") + return False + def _update_inference_buttons(self) -> None: preview_running = self.multi_camera_controller.is_running() self.start_inference_button.setEnabled(preview_running and not self._dlc_active) @@ -2405,11 +2651,14 @@ def _update_dlc_controls_enabled(self) -> None: for widget in widgets: widget.setEnabled(allow_changes) + dlc_controls_enabled = allow_changes and self._backend_name == "dlc" for widget in processor_widgets: - widget.setEnabled(allow_changes) + widget.setEnabled(dlc_controls_enabled) - if hasattr(self, "use_custom_proc_checkbox"): - self.use_custom_proc_checkbox.setEnabled(allow_changes) + self.use_custom_proc_checkbox.setEnabled(dlc_controls_enabled) + + self.action_backend_dlc.setEnabled(allow_changes) + self.action_backend_poet.setEnabled(allow_changes) def _update_camera_controls_enabled(self) -> None: multi_cam_recording = self._rec_manager.is_active @@ -2492,11 +2741,11 @@ def _update_metrics(self) -> None: # --- DLC processor stats --- if hasattr(self, "dlc_stats_label"): if self._dlc_active and self._dlc_initialized: - stats = self._dlc.get_stats() + stats = self._active_pose_processor.get_stats() summary = format_dlc_stats(stats) self.dlc_stats_label.setText(summary) else: - self.dlc_stats_label.setText("DLC processor idle") + self.dlc_stats_label.setText("Processor idle") # Update processor status (connection and recording state) if hasattr(self, "processor_status_label") and self._custom_processor_enabled(): @@ -2513,6 +2762,10 @@ def _update_metrics(self) -> None: def _update_processor_status(self) -> None: """Update processor connection and recording status, handle auto-recording.""" + if self._backend_name != "dlc": + self.processor_status_label.setText("Unavailable for this backend") + return + if not self._custom_processor_enabled(): self.processor_status_label.setText("Disabled") return @@ -2582,27 +2835,46 @@ def _start_inference(self) -> None: if not self.multi_camera_controller.is_running(): self._show_error("Start the camera preview before running pose inference.") return - if not self._configure_dlc(): + if not self._configure_pose_backend(): self._update_inference_buttons() return - self._dlc.reset() + + self._active_pose_processor.reset() self._last_pose = None self._overlay_renderer.clear_runtime_state() + self._dlc_active = True self._dlc_initialized = False - # Update button to show initializing state - self.start_inference_button.setText("Initializing DLCLive!") + backend_name = self._backend_display_name() + + self.start_inference_button.setText(f"Initializing {backend_name}...") self.start_inference_button.setStyleSheet("background-color: #4A90E2; color: white;") self.start_inference_button.setEnabled(False) self.stop_inference_button.setEnabled(True) - self.statusBar().showMessage("Initializing DLCLive…", 3000) + self.statusBar().showMessage( + f"Initializing {backend_name}...", + 3000, + ) self._update_camera_controls_enabled() self._update_dlc_controls_enabled() + def _reset_active_pose_processor( + self, + *, + reset_processor_plugin: bool, + ) -> None: + if self._backend_name == "dlc": + self._dlc.reset( + reset_processor_plugin=reset_processor_plugin, + ) + return + + self._active_pose_processor.reset() + def _stop_inference(self, show_message: bool = True) -> None: - if self._rec_manager.is_active: + if self._rec_manager.is_active and self._backend_name == "dlc": answer = QMessageBox.question( self, "Stop inference while recording?", @@ -2619,7 +2891,7 @@ def _stop_inference(self, show_message: bool = True) -> None: self._dlc_active = False self._dlc_initialized = False # Does NOT invoke the normal rec-stop/save hooks. Persistence is processor-dependent. - self._dlc.reset(reset_processor_plugin=True) + self._reset_active_pose_processor(reset_processor_plugin=True) self._last_pose = None self._overlay_renderer.clear_runtime_state() self._last_processor_vid_recording = False @@ -2704,7 +2976,7 @@ def _on_pose_ready( self._dlc_timing.maybe_log() - def _on_dlc_error(self, message: str) -> None: + def _on_pose_error(self, message: str) -> None: self._stop_inference(show_message=False) self._show_error(message) @@ -2756,24 +3028,29 @@ def _on_bbox_changed(self, _value: int = 0) -> None: if self._current_frame is not None: self._display_frame(self._current_frame, force=True) - def _on_dlc_initialised(self, success: bool) -> None: + def _on_pose_processor_initialised(self, success: bool) -> None: if success: self._dlc_initialized = True - if self._processor_recording_context is not None and not self._processor_recording_started_notified: + if ( + self._backend_name == "dlc" + and self._processor_recording_context is not None + and not self._processor_recording_started_notified + ): self._notify_processor_recording_started(self._processor_recording_context) # Update button to show running state self.start_inference_button.setText("DLCLive running!") self.start_inference_button.setStyleSheet("background-color: #4CAF50; color: white;") self.statusBar().showMessage("DLCLive initialized successfully", 3000) - else: - self._dlc_initialized = False - # Reset button on failure - self.start_inference_button.setText("Start pose inference") - self.start_inference_button.setStyleSheet("") - self.statusBar().showMessage("DLCLive initialization failed", 5000) - # Stop inference since initialization failed - self._stop_inference(show_message=False) + return + + self._dlc_initialized = False + # Reset button on failure + self.start_inference_button.setText("Start pose inference") + self.start_inference_button.setStyleSheet("") + self.statusBar().showMessage("DLCLive initialization failed", 5000) + # Stop inference since initialization failed + self._stop_inference(show_message=False) # ------------------------------------------------------------------ # Helpers @@ -2817,6 +3094,16 @@ def _set_combo_from_color(self, bgr: tuple[int, int, int]) -> None: # ------------------------------------------------------------------ # Qt overrides + def _shutdown_pose_processors(self) -> None: + for backend, processor in tuple(self._pose_processors.items()): + try: + processor.shutdown() + except Exception: + logger.exception( + "Failed to shut down %s pose processor.", + backend, + ) + def closeEvent(self, event: QCloseEvent) -> None: # pragma: no cover - GUI behaviour if self.multi_camera_controller.is_running(): self.multi_camera_controller.stop(wait=True) @@ -2843,7 +3130,7 @@ def closeEvent(self, event: QCloseEvent) -> None: # pragma: no cover - GUI beha pass self._cam_dialog = None - self._dlc.shutdown() + self._shutdown_pose_processors() if hasattr(self, "_metrics_timer"): self._metrics_timer.stop() From 54970f330769d46cd71b075740b89835c30ce6ee Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 14:38:08 +0200 Subject: [PATCH 07/23] Fix tests --- tests/custom_processors/test_processor_rec_context.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/custom_processors/test_processor_rec_context.py b/tests/custom_processors/test_processor_rec_context.py index 85d840f3..caf33b7f 100644 --- a/tests/custom_processors/test_processor_rec_context.py +++ b/tests/custom_processors/test_processor_rec_context.py @@ -252,7 +252,7 @@ def test_main_window_get_processor_instance_prefers_direct_processor(main_window processor = HookProcessor() win = make_window_shell(main_window_cls, processor=processor) - assert win._get_dlc_processor_instance() is processor + assert win._get_dlc_custom_processor_instance() is processor def test_main_window_get_processor_instance_falls_back_to_dlclive_processor(main_window_cls): @@ -260,7 +260,7 @@ def test_main_window_get_processor_instance_falls_back_to_dlclive_processor(main win = main_window_cls.__new__(main_window_cls) win._dlc = SimpleNamespace(_processor=None, _dlc=SimpleNamespace(processor=processor)) - assert win._get_dlc_processor_instance() is processor + assert win._get_dlc_custom_processor_instance() is processor def test_main_window_notify_processor_recording_started_calls_hook(main_window_cls, tmp_path): From 50bf14e159ec79f2b1b63a2ff76b67fe0e9b8179 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:07:29 +0200 Subject: [PATCH 08/23] Add POET weights dialog and backend path memory Introduces a new non-modal POET weights dialog that lets users browse existing .pt/.pth files or download the default checkpoint with progress and error handling. Main window integration adds a "Manage POET weights" action, applies selected POET weights to the model path, and persists/restores both backend choice and backend-specific model paths via QSettings so switching between DLC and POET keeps the correct path. The POET default weights directory was also changed to `Path.cwd() / "TEMP_WEIGHTS/poet"`. --- dlclivegui/gui/main_window.py | 131 ++++++++++++- dlclivegui/gui/misc/weights_dialog.py | 179 ++++++++++++++++++ .../services/inference/models/poet/weights.py | 4 +- 3 files changed, 303 insertions(+), 11 deletions(-) create mode 100644 dlclivegui/gui/misc/weights_dialog.py diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index 0a2325c9..7a80f06f 100644 --- a/dlclivegui/gui/main_window.py +++ b/dlclivegui/gui/main_window.py @@ -9,7 +9,7 @@ import threading import time from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np from PySide6.QtCore import QRect, QSettings, Qt, QTimer, QUrl, Signal @@ -95,6 +95,7 @@ from .misc import layouts as lyts from .misc.drag_spinbox import ScrubSpinBox from .misc.eliding_label import ElidingPathLabel +from .misc.weights_dialog import PoetWeightsDialog from .qt_display.utils import frame_to_pixmap from .theme import LOGO, LOGO_ALPHA, AppStyle, apply_theme @@ -192,6 +193,7 @@ def __init__(self, config: ApplicationSettings | None = None): # UI elements self._current_style: AppStyle = AppStyle.DARK self._cam_dialog: CameraConfigDialog | None = None + self._poet_weights_dialog: PoetWeightsDialog | None = None # Visualization settings (will be updated from config) self._p_cutoff = 0.6 @@ -218,6 +220,7 @@ def __init__(self, config: ApplicationSettings | None = None): self._preview_pixmap = QPixmap(LOGO_ALPHA) self._setup_ui() self._connect_signals() + self._restore_pose_backend() self._apply_config(self._config, restore_local_prefs=True) self._refresh_processors() # Scan and populate processor dropdown self._update_inference_buttons() @@ -428,6 +431,59 @@ def _build_stats_layout(self, stats_widget: QWidget) -> QGridLayout: stats_widget.setLayout(stats_layout) + def _action_manage_poet_weights(self) -> None: + if self._poet_weights_dialog is None: + dialog = PoetWeightsDialog( + self, + initial_path=(self._get_last_poet_weights_path()), + ) + dialog.weights_selected.connect(self._on_poet_weights_selected) + dialog.finished.connect(self._clear_poet_weights_dialog) + + self._poet_weights_dialog = dialog + + self._poet_weights_dialog.show() + self._poet_weights_dialog.raise_() + self._poet_weights_dialog.activateWindow() + + def _on_poet_weights_selected( + self, + path: str, + ) -> None: + self._set_last_poet_weights_path(path) + self._set_backend("poet") + self.model_path_edit.setText(path) + + self.statusBar().showMessage( + f"POET weights selected: {path}", + 5000, + ) + + def _clear_poet_weights_dialog( + self, + _result: int, + ) -> None: + self._poet_weights_dialog = None + + def _set_last_poet_weights_path( + self, + path: str, + ) -> None: + self.settings.setValue( + "poet/weights_path", + path, + ) + + def _get_last_poet_weights_path(self) -> str: + return ( + self.settings.value( + "poet/weights_path", + "", + type=str, + ) + or "" + ) + def _build_menus(self) -> None: # File menu file_menu = self.menuBar().addMenu("&File") @@ -497,6 +553,14 @@ def _build_menus(self) -> None: self._backend_action_group.setExclusive(True) self._backend_action_group.addAction(self.action_backend_dlc) self._backend_action_group.addAction(self.action_backend_poet) + # ------------------------- + models_menu.addSeparator() + self.action_manage_poet_weights = QAction( + "Manage POET weights...", + self, + ) + self.action_manage_poet_weights.triggered.connect(self._action_manage_poet_weights) + models_menu.addAction(self.action_manage_poet_weights) backend_menu.addAction(self.action_backend_dlc) backend_menu.addAction(self.action_backend_poet) @@ -1015,7 +1079,10 @@ def _apply_config(self, config: ApplicationSettings, *, restore_local_prefs: boo # Set DLC settings from config dlc = config.dlc - resolved_model_path = self._model_path_store.resolve(dlc.model_path) + if self._backend_name == "poet": + resolved_model_path = self._get_last_poet_weights_path() + else: + resolved_model_path = self._model_path_store.resolve(dlc.model_path) self.model_path_edit.setText(resolved_model_path) # self.additional_options_edit.setPlainText(json.dumps(dlc.additional_options, indent=2)) @@ -1439,7 +1506,10 @@ def _action_browse_model(self) -> None: selected_path = str(file_path) self.model_path_edit.setText(selected_path) - self._model_path_store.save_if_valid(selected_path) + if self._backend_name == "poet": + self._set_last_poet_weights_path(selected_path) + else: + self._model_path_store.save_if_valid(selected_path) def _action_browse_directory(self) -> None: directory = QFileDialog.getExistingDirectory(self, "Select output directory", str(Path.home())) @@ -2143,18 +2213,40 @@ def _get_pose_processor( return processor + def _remember_current_backend_path(self) -> None: + path = self.model_path_edit.text().strip() + + if not path: + return + + if self._backend_name == "poet": + self._set_last_poet_weights_path(path) + else: + self._model_path_store.save_if_valid(path) + + def _restore_current_backend_path(self) -> None: + if self._backend_name == "poet": + path = self._get_last_poet_weights_path() + else: + path = self._model_path_store.resolve(self._config.dlc.model_path) + + self.model_path_edit.setText(path) + def _set_backend( self, name: str, ) -> None: - if name not in {"dlc", "poet"}: - raise ValueError(f"Unsupported pose backend: {name!r}.") - backend: PoseBackendName = name + normalized_name = name.strip().lower() + if normalized_name not in {"dlc", "poet"}: + raise ValueError(f"Unsupported pose backend: {normalized_name!r}.") + + backend = cast(PoseBackendName, normalized_name) if backend == self._backend_name: self._sync_backend_actions() self._update_backend_ui() + self._restore_current_backend_path() return if self._dlc_active: @@ -2162,11 +2254,10 @@ def _set_backend( self._sync_backend_actions() return + self._remember_current_backend_path() self._disconnect_pose_processor_signals() self._backend_name = backend - - # Construct the processor lazily before connecting its signals. self._get_pose_processor(backend) self._last_pose = None @@ -2175,6 +2266,7 @@ def _set_backend( self._connect_pose_processor_signals() self._sync_backend_actions() self._update_backend_ui() + self._restore_current_backend_path() self.settings.setValue( "app/backend", @@ -2186,6 +2278,22 @@ def _set_backend( 3000, ) + def _restore_pose_backend(self) -> None: + saved_backend = self.settings.value( + "app/backend", + "dlc", + type=str, + ) + + if saved_backend not in {"dlc", "poet"}: + logger.warning( + "Ignoring unsupported saved pose backend: %r", + saved_backend, + ) + saved_backend = "dlc" + + self._set_backend(saved_backend) + def _backend_display_name(self) -> str: if self._backend_name == "poet": return "POET" @@ -2659,6 +2767,7 @@ def _update_dlc_controls_enabled(self) -> None: self.action_backend_dlc.setEnabled(allow_changes) self.action_backend_poet.setEnabled(allow_changes) + self.action_manage_poet_weights.setEnabled(allow_changes) def _update_camera_controls_enabled(self) -> None: multi_cam_recording = self._rec_manager.is_active @@ -3138,7 +3247,11 @@ def closeEvent(self, event: QCloseEvent) -> None: # pragma: no cover - GUI beha self._display_timer.stop() # Remember model path on exit - self._model_path_store.save_if_valid(self.model_path_edit.text().strip()) + current_model_path = self.model_path_edit.text().strip() + if self._backend_name == "poet": + self._set_last_poet_weights_path(current_model_path) + else: + self._model_path_store.save_if_valid(current_model_path) # Remember processor folder on exit if hasattr(self, "processor_folder_edit"): diff --git a/dlclivegui/gui/misc/weights_dialog.py b/dlclivegui/gui/misc/weights_dialog.py new file mode 100644 index 00000000..db231fb1 --- /dev/null +++ b/dlclivegui/gui/misc/weights_dialog.py @@ -0,0 +1,179 @@ +"""Dialog for selecting or downloading POET weights.""" + +from __future__ import annotations + +from pathlib import Path + +from PySide6.QtCore import QThread, Signal, Slot +from PySide6.QtWidgets import ( + QDialog, + QFileDialog, + QHBoxLayout, + QLabel, + QMessageBox, + QProgressBar, + QPushButton, + QVBoxLayout, +) + +from dlclivegui.services.inference.models.poet.weights import ( + POET_WEIGHTS_FILENAME, + POET_WEIGHTS_URL, + WeightsDownloadWorker, + poet_default_weights_dir, +) + + +class PoetWeightsDialog(QDialog): + """Select existing POET weights or download the default checkpoint.""" + + weights_selected = Signal(str) + + def __init__( + self, + parent=None, + *, + initial_path: str = "", + ) -> None: + super().__init__(parent) + self.setWindowTitle("POET weights") + self.setModal(False) + + self._selected_path = initial_path + self._download_thread: QThread | None = None + self._download_worker: WeightsDownloadWorker | None = None + + self.path_label = QLabel(initial_path or "No weights selected") + self.path_label.setWordWrap(True) + + self.progress_bar = QProgressBar() + self.progress_bar.setRange(0, 100) + self.progress_bar.setValue(0) + self.progress_bar.setVisible(False) + + self.browse_button = QPushButton("Select existing weights...") + self.download_button = QPushButton("Download default weights") + self.use_button = QPushButton("Use selected weights") + self.use_button.setEnabled(self._is_valid_weights_path(initial_path)) + + button_row = QHBoxLayout() + button_row.addWidget(self.browse_button) + button_row.addWidget(self.download_button) + + layout = QVBoxLayout(self) + layout.addWidget(self.path_label) + layout.addWidget(self.progress_bar) + layout.addLayout(button_row) + layout.addWidget(self.use_button) + + self.browse_button.clicked.connect(self._browse) + self.download_button.clicked.connect(self._download) + self.use_button.clicked.connect(self._use_selected) + + @staticmethod + def _is_valid_weights_path( + path: str, + ) -> bool: + candidate = Path(path).expanduser() + + return candidate.is_file() and candidate.suffix.lower() in { + ".pt", + ".pth", + } + + @Slot() + def _browse(self) -> None: + start_path = ( + self._selected_path if self._is_valid_weights_path(self._selected_path) else str(poet_default_weights_dir()) + ) + + path, _ = QFileDialog.getOpenFileName( + self, + "Select POET weights", + start_path, + "POET weights (*.pt *.pth)", + ) + + if path: + self._set_selected_path(path) + + @Slot() + def _download(self) -> None: + if self._download_thread is not None: + return + + destination = poet_default_weights_dir() / POET_WEIGHTS_FILENAME + + thread = QThread(self) + worker = WeightsDownloadWorker( + POET_WEIGHTS_URL, + destination, + ) + worker.moveToThread(thread) + + thread.started.connect(worker.run) + worker.progress.connect(self.progress_bar.setValue) + worker.finished.connect(self._download_finished) + worker.error.connect(self._download_failed) + + worker.finished.connect(thread.quit) + worker.error.connect(thread.quit) + + thread.finished.connect(self._download_cleanup) + thread.finished.connect(worker.deleteLater) + thread.finished.connect(thread.deleteLater) + + self._download_thread = thread + self._download_worker = worker + + self.progress_bar.setVisible(True) + self.progress_bar.setValue(0) + self.browse_button.setEnabled(False) + self.download_button.setEnabled(False) + self.use_button.setEnabled(False) + + thread.start() + + @Slot(str) + def _download_finished( + self, + path: str, + ) -> None: + self.progress_bar.setValue(100) + self._set_selected_path(path) + + @Slot(str) + def _download_failed( + self, + message: str, + ) -> None: + QMessageBox.critical( + self, + "POET weights download failed", + message, + ) + + @Slot() + def _download_cleanup(self) -> None: + self._download_thread = None + self._download_worker = None + + self.browse_button.setEnabled(True) + self.download_button.setEnabled(True) + self.use_button.setEnabled(self._is_valid_weights_path(self._selected_path)) + + def _set_selected_path( + self, + path: str, + ) -> None: + self._selected_path = path + self.path_label.setText(path) + self.use_button.setEnabled(self._is_valid_weights_path(path)) + + @Slot() + def _use_selected(self) -> None: + if not self._is_valid_weights_path(self._selected_path): + return + + self.weights_selected.emit(self._selected_path) + self.accept() diff --git a/dlclivegui/services/inference/models/poet/weights.py b/dlclivegui/services/inference/models/poet/weights.py index d361ba83..f1b2b90a 100644 --- a/dlclivegui/services/inference/models/poet/weights.py +++ b/dlclivegui/services/inference/models/poet/weights.py @@ -2,7 +2,6 @@ import logging import os -import tempfile import urllib.request from pathlib import Path @@ -17,7 +16,8 @@ def poet_default_weights_dir() -> Path: # TODO: move to an app cache rather than a temporary directory - return Path(tempfile.gettempdir()) / "dlclivegui" / "poet" + # return Path(tempfile.gettempdir()) / "dlclivegui" / "poet" + return Path.cwd() / "TEMP_WEIGHTS/poet" class WeightsDownloadWorker(QObject): From 3fad762afc4a83907010dec9aebbce338d0b624d Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:15:35 +0200 Subject: [PATCH 09/23] Add tests for inference pose processor Introduce a new `test_processor.py` suite for inference processing. The tests cover processor configuration, pose signal emission with packets, worker startup on first frame, initialization failure handling, queue-drop accounting, reset behavior (including stop-worker failure), and backend factory creation paths for DLC/POET plus unknown-backend rejection. --- tests/services/inference/test_processor.py | 250 +++++++++++++++++++++ 1 file changed, 250 insertions(+) create mode 100644 tests/services/inference/test_processor.py diff --git a/tests/services/inference/test_processor.py b/tests/services/inference/test_processor.py new file mode 100644 index 00000000..fcaa7dea --- /dev/null +++ b/tests/services/inference/test_processor.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import queue +import threading +import time + +import numpy as np +import pytest + +from dlclivegui.services.dlc_processor import DLCLiveProcessor +from dlclivegui.services.inference.base import ( + PoseBackend, + PoseBackends, + PosePacket, + PoseSource, +) +from dlclivegui.services.inference.processor import PoseProcessor, create_pose_processor + + +class FakePoseBackend(PoseBackend): + def __init__( + self, + *, + pose: np.ndarray | None = None, + initialization_error: Exception | None = None, + inference_error: Exception | None = None, + ) -> None: + self.pose = pose + self.initialization_error = initialization_error + self.inference_error = inference_error + + self.initialized_with: np.ndarray | None = None + self.received_frames: list[np.ndarray] = [] + self.closed = False + + def init_inference( + self, + init_frame: np.ndarray, + ) -> None: + if self.initialization_error is not None: + raise self.initialization_error + + self.initialized_with = init_frame.copy() + + def get_pose( + self, + frame: np.ndarray, + frame_time: float | None = None, + ) -> np.ndarray | None: + del frame_time + + if self.inference_error is not None: + raise self.inference_error + + self.received_frames.append(frame.copy()) + return self.pose + + def make_pose_packet( + self, + pose: np.ndarray | None, + ) -> PosePacket: + return PosePacket( + schema_version=1, + keypoints=pose, + keypoint_names=["a", "b"], + individual_ids=None, + skeleton_id="test.line", + skeleton_edges=(("a", "b"),), + source=PoseSource( + backend=PoseBackends.POET, + model_type=None, + ), + raw=pose, + ) + + def close(self) -> None: + self.closed = True + + +def test_configure_stores_backend_factory() -> None: + processor = PoseProcessor() + + processor.configure(FakePoseBackend) + + assert processor.is_configured() + + +def test_process_frame_emits_pose_result_with_packet( + qtbot, +) -> None: + pose = np.array( + [ + [1.0, 2.0, 0.9], + [3.0, 4.0, 0.8], + ], + dtype=np.float32, + ) + backend = FakePoseBackend(pose=pose) + + processor = PoseProcessor() + processor._backend = backend + + with qtbot.waitSignal( + processor.pose_ready, + timeout=1000, + ) as blocker: + processor._process_frame( + np.zeros((10, 10, 3), dtype=np.uint8), + timestamp=10.0, + enqueue_time=time.perf_counter(), + queue_wait_time=0.0, + ) + + result = blocker.args[0] + + assert result.timestamp == 10.0 + assert result.pose is pose + assert result.packet.skeleton_id == "test.line" + assert result.packet.skeleton_edges == (("a", "b"),) + + +def test_enqueue_first_frame_starts_worker( + qtbot, +) -> None: + backend = FakePoseBackend(pose=np.zeros((2, 3), dtype=np.float32)) + processor = PoseProcessor() + processor.configure(lambda: backend) + + frame = np.zeros((10, 10, 3), dtype=np.uint8) + + with qtbot.waitSignal( + processor.initialized, + timeout=2000, + ) as blocker: + processor.enqueue_frame(frame, timestamp=1.0) + + assert blocker.args == [True] + assert backend.initialized_with is not None + + processor.shutdown() + assert backend.closed + + +def test_initialization_failure_emits_error_and_false( + qtbot, +) -> None: + processor = PoseProcessor() + processor.configure(lambda: FakePoseBackend(initialization_error=RuntimeError("broken backend"))) + + errors: list[str] = [] + initialized: list[bool] = [] + + processor.error.connect(errors.append) + processor.initialized.connect(initialized.append) + + processor.enqueue_frame( + np.zeros((10, 10, 3), dtype=np.uint8), + timestamp=1.0, + ) + + qtbot.waitUntil( + lambda: bool(initialized), + timeout=2000, + ) + + assert errors == ["broken backend"] + assert initialized == [False] + assert processor._backend is None + assert processor._worker_thread is None + assert processor._queue is None + + +def test_enqueue_frame_counts_drop_when_queue_is_full() -> None: + processor = PoseProcessor() + processor._worker_thread = threading.current_thread() + + processor._queue = queue.Queue(maxsize=1) + processor._queue.put_nowait( + ( + np.zeros((1, 1, 3), dtype=np.uint8), + 1.0, + 1.0, + ) + ) + + processor.enqueue_frame( + np.zeros((1, 1, 3), dtype=np.uint8), + timestamp=2.0, + ) + + stats = processor.get_stats() + + assert stats.frames_dropped == 1 + + +def test_reset_clears_statistics() -> None: + processor = PoseProcessor() + processor._frames_enqueued = 4 + processor._frames_processed = 3 + processor._frames_dropped = 2 + processor._latencies.append(1.0) + processor._processing_times.extend([1.0, 2.0]) + + processor.reset() + + stats = processor.get_stats() + + assert stats.frames_enqueued == 0 + assert stats.frames_processed == 0 + assert stats.frames_dropped == 0 + assert stats.average_latency == 0.0 + assert stats.processing_fps == 0.0 + + +def test_reset_raises_when_worker_does_not_stop( + monkeypatch, +) -> None: + processor = PoseProcessor() + + monkeypatch.setattr( + processor, + "_stop_worker", + lambda: False, + ) + + with pytest.raises( + RuntimeError, + match="Failed to stop worker thread", + ): + processor.reset() + + +def test_create_dlc_pose_processor() -> None: + processor = create_pose_processor("dlc") + + assert isinstance(processor, DLCLiveProcessor) + + +def test_create_poet_pose_processor() -> None: + processor = create_pose_processor("poet") + + assert isinstance(processor, PoseProcessor) + + +def test_create_pose_processor_rejects_unknown_backend() -> None: + with pytest.raises( + ValueError, + match="Unsupported pose backend", + ): + create_pose_processor("unknown") From 973bdd0826d227295800afa2349e5e35675e6eca Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:15:49 +0200 Subject: [PATCH 10/23] Add main window backend switching tests Introduce GUI tests for backend handling in the main window, focused on the POET integration. The new coverage verifies switching UI/state updates, lazy POET processor creation, blocking backend changes during active inference, per-backend model path restoration, restoring POET from saved settings, and successful POET configuration/processor setup. --- tests/gui/main_window/test_backends.py | 107 +++++++++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 tests/gui/main_window/test_backends.py diff --git a/tests/gui/main_window/test_backends.py b/tests/gui/main_window/test_backends.py new file mode 100644 index 00000000..6ad82a02 --- /dev/null +++ b/tests/gui/main_window/test_backends.py @@ -0,0 +1,107 @@ +from dlclivegui.services.inference.processor import PoseProcessor + + +def test_set_backend_switches_to_poet( + window, +) -> None: + window._set_backend("poet") + + assert window._backend_name == "poet" + assert window.action_backend_poet.isChecked() + assert not window.action_backend_dlc.isChecked() + assert window.inference_group.title() == "POET" + assert "poet_weights" in (window.model_path_edit.placeholderText()) + + +def test_poet_processor_is_created_lazily( + window, +) -> None: + assert "poet" not in window._pose_processors + + window._set_backend("poet") + + assert "poet" in window._pose_processors + assert isinstance( + window._pose_processors["poet"], + PoseProcessor, + ) + + +def test_backend_switch_is_rejected_during_inference( + window, + monkeypatch, +) -> None: + warnings: list[str] = [] + + window._dlc_active = True + monkeypatch.setattr( + window, + "_show_warning", + warnings.append, + ) + + window._set_backend("poet") + + assert window._backend_name == "dlc" + assert warnings == ["Stop pose inference before switching backends."] + + +def test_switching_backends_restores_separate_paths( + window, + tmp_path, + monkeypatch, +) -> None: + dlc_path = tmp_path / "dlc_model.pth" + poet_path = tmp_path / "poet_weights.pth" + + dlc_path.touch() + poet_path.touch() + + monkeypatch.setattr( + window._model_path_store, + "resolve", + lambda _configured: str(dlc_path), + ) + + window.model_path_edit.setText(str(dlc_path)) + window._set_last_poet_weights_path(str(poet_path)) + + window._set_backend("poet") + assert window.model_path_edit.text() == str(poet_path) + + window._set_backend("dlc") + assert window.model_path_edit.text() == str(dlc_path) + + +def test_restore_pose_backend_restores_poet( + window, +) -> None: + window.settings.setValue( + "app/backend", + "poet", + ) + + window._restore_pose_backend() + + assert window._backend_name == "poet" + assert window.action_backend_poet.isChecked() + + +def test_configure_poet_registers_backend_factory( + window, + tmp_path, +) -> None: + checkpoint = tmp_path / "weights.pth" + checkpoint.touch() + + window._set_backend("poet") + window.model_path_edit.setText(str(checkpoint)) + + result = window._configure_poet() + + processor = window._active_pose_processor + + assert result is True + assert isinstance(processor, PoseProcessor) + assert processor.is_configured() + assert processor._backend is None From 6b62a7847e7fd9384c9368b40639f630bfc75eab Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:16:02 +0200 Subject: [PATCH 11/23] Add UI tests for POET weights dialog Introduce a new test module covering `PoetWeightsDialog` behavior: default state without an initial path, handling a valid initial path, emitting `weights_selected` on use, updating state after browse selection, and restoring button states during download cleanup. --- tests/gui/ui_blocks/test_weights_dialog.py | 98 ++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 tests/gui/ui_blocks/test_weights_dialog.py diff --git a/tests/gui/ui_blocks/test_weights_dialog.py b/tests/gui/ui_blocks/test_weights_dialog.py new file mode 100644 index 00000000..c0812823 --- /dev/null +++ b/tests/gui/ui_blocks/test_weights_dialog.py @@ -0,0 +1,98 @@ +from PySide6.QtCore import Qt +from PySide6.QtWidgets import QFileDialog + +from dlclivegui.gui.misc.weights_dialog import PoetWeightsDialog + + +def test_dialog_without_initial_path_disables_use_button( + qtbot, +) -> None: + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + assert dialog.path_label.text() == "No weights selected" + assert not dialog.use_button.isEnabled() + assert not dialog.progress_bar.isVisible() + + +def test_dialog_accepts_valid_initial_path( + qtbot, + tmp_path, +) -> None: + path = tmp_path / "weights.pth" + path.touch() + + dialog = PoetWeightsDialog(initial_path=str(path)) + qtbot.addWidget(dialog) + + assert dialog.path_label.text() == str(path) + assert dialog.use_button.isEnabled() + + +def test_use_selected_emits_path( + qtbot, + tmp_path, +) -> None: + path = tmp_path / "weights.pth" + path.touch() + + dialog = PoetWeightsDialog(initial_path=str(path)) + qtbot.addWidget(dialog) + + with qtbot.waitSignal( + dialog.weights_selected, + timeout=1000, + ) as blocker: + qtbot.mouseClick( + dialog.use_button, + Qt.MouseButton.LeftButton, + ) + + assert blocker.args == [str(path)] + + +def test_browse_sets_selected_path( + qtbot, + monkeypatch, + tmp_path, +) -> None: + path = tmp_path / "weights.pth" + path.touch() + + monkeypatch.setattr( + QFileDialog, + "getOpenFileName", + lambda *args, **kwargs: ( + str(path), + "POET weights (*.pt *.pth)", + ), + ) + + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + dialog._browse() + + assert dialog.path_label.text() == str(path) + assert dialog.use_button.isEnabled() + + +def test_download_cleanup_restores_controls( + qtbot, + tmp_path, +) -> None: + path = tmp_path / "weights.pth" + path.touch() + + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + dialog._set_selected_path(str(path)) + dialog.browse_button.setEnabled(False) + dialog.download_button.setEnabled(False) + + dialog._download_cleanup() + + assert dialog.browse_button.isEnabled() + assert dialog.download_button.isEnabled() + assert dialog.use_button.isEnabled() From b09ac0078f32bc131552a640faf238c1e8ce81e8 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:16:42 +0200 Subject: [PATCH 12/23] Add POET backend and download worker tests Introduces a new `test_poet.py` suite covering POET inference model behavior and weight download handling. The tests validate checkpoint path/extension errors, ensure pose packets include correct POET metadata (including empty-detection cases), and verify `WeightsDownloadWorker` behavior for reusing existing files, writing downloaded content, emitting progress/finished signals, and cleaning up partial files on download failure. --- tests/services/inference/models/test_poet.py | 183 +++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 tests/services/inference/models/test_poet.py diff --git a/tests/services/inference/models/test_poet.py b/tests/services/inference/models/test_poet.py new file mode 100644 index 00000000..a065d711 --- /dev/null +++ b/tests/services/inference/models/test_poet.py @@ -0,0 +1,183 @@ +import urllib.request +from io import BytesIO + +import numpy as np +import pytest + +from dlclivegui.services.inference.base import PoseBackends +from dlclivegui.services.inference.models.poet.poet_processor import ( + POET_KEYPOINT_NAMES, + POET_SKELETON_EDGES, + POET_SKELETON_ID, + POETBackend, +) +from dlclivegui.services.inference.models.poet.weights import WeightsDownloadWorker + + +def test_poet_backend_rejects_missing_checkpoint( + tmp_path, +) -> None: + with pytest.raises( + FileNotFoundError, + match="checkpoint not found", + ): + POETBackend(str(tmp_path / "missing.pth")) + + +def test_poet_backend_rejects_unsupported_extension( + tmp_path, +) -> None: + checkpoint = tmp_path / "weights.txt" + checkpoint.touch() + + with pytest.raises( + ValueError, + match=r"\.pt or \.pth", + ): + POETBackend(str(checkpoint)) + + +def test_make_pose_packet_contains_poet_metadata( + tmp_path, +) -> None: + checkpoint = tmp_path / "weights.pth" + checkpoint.touch() + + backend = POETBackend(str(checkpoint)) + pose = np.zeros((2, 17, 3), dtype=np.float32) + + packet = backend.make_pose_packet(pose) + + assert packet.keypoints is pose + assert packet.keypoint_names == list(POET_KEYPOINT_NAMES) + assert packet.skeleton_id == POET_SKELETON_ID + assert packet.skeleton_edges == POET_SKELETON_EDGES + assert packet.individual_ids == [ + "person_0", + "person_1", + ] + assert packet.source.backend == PoseBackends.POET + + +def test_make_pose_packet_without_detections_keeps_metadata( + tmp_path, +) -> None: + checkpoint = tmp_path / "weights.pth" + checkpoint.touch() + + backend = POETBackend(str(checkpoint)) + + packet = backend.make_pose_packet(None) + + assert packet.keypoints is None + assert packet.individual_ids == [] + assert packet.keypoint_names == list(POET_KEYPOINT_NAMES) + assert packet.skeleton_edges == POET_SKELETON_EDGES + + +def test_download_worker_reuses_existing_file( + qtbot, + tmp_path, +) -> None: + destination = tmp_path / "weights.pth" + destination.write_bytes(b"existing") + + worker = WeightsDownloadWorker( + "https://example.invalid/weights", + destination, + ) + + progress: list[int] = [] + + worker.progress.connect(progress.append) + + with qtbot.waitSignal( + worker.finished, + timeout=1000, + ) as blocker: + worker.run() + + assert blocker.args == [str(destination)] + assert progress == [100] + assert destination.read_bytes() == b"existing" + + +class FakeResponse(BytesIO): + def __init__(self, content: bytes) -> None: + super().__init__(content) + self.length = len(content) + + def __enter__(self): + return self + + def __exit__( + self, + exc_type, + exc_value, + traceback, + ) -> None: + self.close() + + +def test_download_worker_writes_destination( + qtbot, + monkeypatch, + tmp_path, +) -> None: + destination = tmp_path / "weights.pth" + content = b"model-data" + + monkeypatch.setattr( + urllib.request, + "urlopen", + lambda _request: FakeResponse(content), + ) + + worker = WeightsDownloadWorker( + "https://example.invalid/weights", + destination, + ) + + with qtbot.waitSignal( + worker.finished, + timeout=1000, + ) as blocker: + worker.run() + + assert blocker.args == [str(destination)] + assert destination.read_bytes() == content + assert not destination.with_suffix(".pth.part").exists() + + +def test_download_worker_reports_failure_and_cleans_partial_file( + qtbot, + monkeypatch, + tmp_path, +) -> None: + destination = tmp_path / "weights.pth" + partial_path = destination.with_suffix(".pth.part") + + def failing_urlopen(_request): + partial_path.write_bytes(b"partial") + raise OSError("download failed") + + monkeypatch.setattr( + urllib.request, + "urlopen", + failing_urlopen, + ) + + worker = WeightsDownloadWorker( + "https://example.invalid/weights", + destination, + ) + + with qtbot.waitSignal( + worker.error, + timeout=1000, + ) as blocker: + worker.run() + + assert blocker.args == ["download failed"] + assert not partial_path.exists() + assert not destination.exists() From f914fbb2f10979922adc0456950f538f0b0fdce5 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:27:22 +0200 Subject: [PATCH 13/23] Simplify POET weights download dialog Refactor the POET weights dialog into a download-only flow: remove manual file selection/use actions, show a fixed destination path with an eliding copy-friendly label, and auto-emit the downloaded path before closing on success. The dialog now provides clearer status/progress messaging and blocks closing while a download is active. Also ignore `/TEMP_WEIGHTS` in git. --- .gitignore | 1 + dlclivegui/gui/misc/eliding_label.py | 1 + dlclivegui/gui/misc/weights_dialog.py | 141 ++++++++++++-------------- 3 files changed, 68 insertions(+), 75 deletions(-) diff --git a/.gitignore b/.gitignore index 1782ab32..4ec6a236 100644 --- a/.gitignore +++ b/.gitignore @@ -118,3 +118,4 @@ uv.lock profile*.svg scalene*.json scalene*.html +/TEMP_WEIGHTS diff --git a/dlclivegui/gui/misc/eliding_label.py b/dlclivegui/gui/misc/eliding_label.py index ed8597e3..77b4fa8a 100644 --- a/dlclivegui/gui/misc/eliding_label.py +++ b/dlclivegui/gui/misc/eliding_label.py @@ -1,3 +1,4 @@ +# dlclivegui/gui/misc/eliding_label.py from PySide6.QtCore import Qt from PySide6.QtGui import QCursor, QGuiApplication from PySide6.QtWidgets import QLabel, QSizePolicy, QToolTip diff --git a/dlclivegui/gui/misc/weights_dialog.py b/dlclivegui/gui/misc/weights_dialog.py index db231fb1..52f8afcb 100644 --- a/dlclivegui/gui/misc/weights_dialog.py +++ b/dlclivegui/gui/misc/weights_dialog.py @@ -1,13 +1,10 @@ -"""Dialog for selecting or downloading POET weights.""" +"""Dialog for downloading the default POET weights.""" from __future__ import annotations -from pathlib import Path - from PySide6.QtCore import QThread, Signal, Slot from PySide6.QtWidgets import ( QDialog, - QFileDialog, QHBoxLayout, QLabel, QMessageBox, @@ -16,6 +13,7 @@ QVBoxLayout, ) +from dlclivegui.gui.misc.eliding_label import ElidingPathLabel from dlclivegui.services.inference.models.poet.weights import ( POET_WEIGHTS_FILENAME, POET_WEIGHTS_URL, @@ -25,94 +23,77 @@ class PoetWeightsDialog(QDialog): - """Select existing POET weights or download the default checkpoint.""" + """Download the default POET checkpoint and return its path.""" - weights_selected = Signal(str) + weights_downloaded = Signal(str) - def __init__( - self, - parent=None, - *, - initial_path: str = "", - ) -> None: + def __init__(self, parent=None) -> None: super().__init__(parent) - self.setWindowTitle("POET weights") + self.setWindowTitle("Download POET weights") self.setModal(False) + self.setMinimumWidth(480) - self._selected_path = initial_path + self._destination = poet_default_weights_dir() / POET_WEIGHTS_FILENAME self._download_thread: QThread | None = None self._download_worker: WeightsDownloadWorker | None = None + self._completed_path: str | None = None + + description = QLabel( + "Download the default POET checkpoint. When the download " + "finishes, the model will be selected automatically and " + "this window will close." + ) + description.setWordWrap(True) - self.path_label = QLabel(initial_path or "No weights selected") - self.path_label.setWordWrap(True) + destination_title = QLabel("Destination:") + + self.path_label = ElidingPathLabel(str(self._destination)) + self.path_label.setToolTip(f"Click to copy:\n{self._destination}") + + self.status_label = QLabel("Ready to download.") + self.status_label.setWordWrap(True) self.progress_bar = QProgressBar() self.progress_bar.setRange(0, 100) self.progress_bar.setValue(0) - self.progress_bar.setVisible(False) + self.progress_bar.setTextVisible(True) - self.browse_button = QPushButton("Select existing weights...") - self.download_button = QPushButton("Download default weights") - self.use_button = QPushButton("Use selected weights") - self.use_button.setEnabled(self._is_valid_weights_path(initial_path)) + self.download_button = QPushButton("Download weights") + self.close_button = QPushButton("Close") button_row = QHBoxLayout() - button_row.addWidget(self.browse_button) + button_row.addStretch(1) + button_row.addWidget(self.close_button) button_row.addWidget(self.download_button) layout = QVBoxLayout(self) + layout.addWidget(description) + layout.addSpacing(6) + layout.addWidget(destination_title) layout.addWidget(self.path_label) + layout.addSpacing(6) + layout.addWidget(self.status_label) layout.addWidget(self.progress_bar) layout.addLayout(button_row) - layout.addWidget(self.use_button) - self.browse_button.clicked.connect(self._browse) self.download_button.clicked.connect(self._download) - self.use_button.clicked.connect(self._use_selected) - - @staticmethod - def _is_valid_weights_path( - path: str, - ) -> bool: - candidate = Path(path).expanduser() - - return candidate.is_file() and candidate.suffix.lower() in { - ".pt", - ".pth", - } - - @Slot() - def _browse(self) -> None: - start_path = ( - self._selected_path if self._is_valid_weights_path(self._selected_path) else str(poet_default_weights_dir()) - ) - - path, _ = QFileDialog.getOpenFileName( - self, - "Select POET weights", - start_path, - "POET weights (*.pt *.pth)", - ) - - if path: - self._set_selected_path(path) + self.close_button.clicked.connect(self.reject) @Slot() def _download(self) -> None: if self._download_thread is not None: return - destination = poet_default_weights_dir() / POET_WEIGHTS_FILENAME - thread = QThread(self) worker = WeightsDownloadWorker( POET_WEIGHTS_URL, - destination, + self._destination, ) worker.moveToThread(thread) thread.started.connect(worker.run) - worker.progress.connect(self.progress_bar.setValue) + + worker.progress.connect(self._update_progress) worker.finished.connect(self._download_finished) worker.error.connect(self._download_failed) @@ -125,28 +106,41 @@ def _download(self) -> None: self._download_thread = thread self._download_worker = worker + self._completed_path = None - self.progress_bar.setVisible(True) + self.status_label.setText("Downloading POET weights...") self.progress_bar.setValue(0) - self.browse_button.setEnabled(False) + self.download_button.setEnabled(False) - self.use_button.setEnabled(False) + self.close_button.setEnabled(False) thread.start() + @Slot(int) + def _update_progress( + self, + value: int, + ) -> None: + self.progress_bar.setValue(value) + self.status_label.setText(f"Downloading POET weights... {value}%") + @Slot(str) def _download_finished( self, path: str, ) -> None: + self._completed_path = path self.progress_bar.setValue(100) - self._set_selected_path(path) + self.status_label.setText("Download complete. Returning to the main window...") @Slot(str) def _download_failed( self, message: str, ) -> None: + self._completed_path = None + self.status_label.setText("Download failed.") + QMessageBox.critical( self, "POET weights download failed", @@ -158,22 +152,19 @@ def _download_cleanup(self) -> None: self._download_thread = None self._download_worker = None - self.browse_button.setEnabled(True) - self.download_button.setEnabled(True) - self.use_button.setEnabled(self._is_valid_weights_path(self._selected_path)) + completed_path = self._completed_path + self._completed_path = None - def _set_selected_path( - self, - path: str, - ) -> None: - self._selected_path = path - self.path_label.setText(path) - self.use_button.setEnabled(self._is_valid_weights_path(path)) + if completed_path is not None: + self.weights_downloaded.emit(completed_path) + self.accept() + return - @Slot() - def _use_selected(self) -> None: - if not self._is_valid_weights_path(self._selected_path): + self.download_button.setEnabled(True) + self.close_button.setEnabled(True) + + def reject(self) -> None: + if self._download_thread is not None: return - self.weights_selected.emit(self._selected_path) - self.accept() + super().reject() From 0329ec618f9c68be83a7eaf9f1f85263bd1b59a5 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:27:38 +0200 Subject: [PATCH 14/23] Auto-apply downloaded POET weights Update the POET weights dialog flow to react to `weights_downloaded` instead of manual selection. Downloaded weights are now persisted as the last path, the backend is switched to POET, the model path field is populated automatically, and a status message is shown. The bounding-box coordinate controls were also repositioned in the form to sit with the bbox settings block for clearer UI grouping. --- dlclivegui/gui/main_window.py | 75 ++++++++++++++++++++--------------- 1 file changed, 43 insertions(+), 32 deletions(-) diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index 7a80f06f..eddbcf6f 100644 --- a/dlclivegui/gui/main_window.py +++ b/dlclivegui/gui/main_window.py @@ -433,11 +433,9 @@ def _build_stats_layout(self, stats_widget: QWidget) -> QGridLayout: def _action_manage_poet_weights(self) -> None: if self._poet_weights_dialog is None: - dialog = PoetWeightsDialog( - self, - initial_path=(self._get_last_poet_weights_path()), - ) - dialog.weights_selected.connect(self._on_poet_weights_selected) + dialog = PoetWeightsDialog(self) + + dialog.weights_downloaded.connect(self._on_poet_weights_downloaded) dialog.finished.connect(self._clear_poet_weights_dialog) self._poet_weights_dialog = dialog @@ -446,6 +444,19 @@ def _action_manage_poet_weights(self) -> None: self._poet_weights_dialog.raise_() self._poet_weights_dialog.activateWindow() + def _on_poet_weights_downloaded( + self, + path: str, + ) -> None: + self._set_last_poet_weights_path(path) + self._set_backend("poet") + self.model_path_edit.setText(path) + + self.statusBar().showMessage( + f"POET weights ready: {path}", + 5000, + ) + def _on_poet_weights_selected( self, path: str, @@ -910,6 +921,33 @@ def _build_viz_group(self) -> QGroupBox: ) form.addRow(bbox_settings) + bbox_layout = QHBoxLayout() + self.bbox_x0_spin = ScrubSpinBox() + self.bbox_x0_spin.setRange(0, 7680) + self.bbox_x0_spin.setPrefix("x0:") + self.bbox_x0_spin.setValue(0) + bbox_layout.addWidget(self.bbox_x0_spin) + + self.bbox_y0_spin = ScrubSpinBox() + self.bbox_y0_spin.setRange(0, 4320) + self.bbox_y0_spin.setPrefix("y0:") + self.bbox_y0_spin.setValue(0) + bbox_layout.addWidget(self.bbox_y0_spin) + + self.bbox_x1_spin = ScrubSpinBox() + self.bbox_x1_spin.setRange(0, 7680) + self.bbox_x1_spin.setPrefix("x1:") + self.bbox_x1_spin.setValue(100) + bbox_layout.addWidget(self.bbox_x1_spin) + + self.bbox_y1_spin = ScrubSpinBox() + self.bbox_y1_spin.setRange(0, 4320) + self.bbox_y1_spin.setPrefix("y1:") + self.bbox_y1_spin.setValue(100) + bbox_layout.addWidget(self.bbox_y1_spin) + + form.addRow("Coordinates", bbox_layout) + # Skeleton overlay self.show_skeleton_checkbox = QCheckBox("Display skeleton") self.show_skeleton_checkbox.setChecked(False) @@ -950,33 +988,6 @@ def _build_viz_group(self) -> QGroupBox: self.skeleton_thickness_spin, ) - bbox_layout = QHBoxLayout() - self.bbox_x0_spin = ScrubSpinBox() - self.bbox_x0_spin.setRange(0, 7680) - self.bbox_x0_spin.setPrefix("x0:") - self.bbox_x0_spin.setValue(0) - bbox_layout.addWidget(self.bbox_x0_spin) - - self.bbox_y0_spin = ScrubSpinBox() - self.bbox_y0_spin.setRange(0, 4320) - self.bbox_y0_spin.setPrefix("y0:") - self.bbox_y0_spin.setValue(0) - bbox_layout.addWidget(self.bbox_y0_spin) - - self.bbox_x1_spin = ScrubSpinBox() - self.bbox_x1_spin.setRange(0, 7680) - self.bbox_x1_spin.setPrefix("x1:") - self.bbox_x1_spin.setValue(100) - bbox_layout.addWidget(self.bbox_x1_spin) - - self.bbox_y1_spin = ScrubSpinBox() - self.bbox_y1_spin.setRange(0, 4320) - self.bbox_y1_spin.setPrefix("y1:") - self.bbox_y1_spin.setValue(100) - bbox_layout.addWidget(self.bbox_y1_spin) - - form.addRow("Coordinates", bbox_layout) - return group # ------------------------------------------------------------------ signals From e3785d95fe35402564b2fadd0ff0be67c5efb14f Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:50:42 +0200 Subject: [PATCH 15/23] Revise POET weights dialog test coverage Reworks `test_weights_dialog.py` to match the updated POET download dialog behavior. The tests now validate default destination/status UI, progress updates, successful download completion and cleanup signaling, failure cleanup control restoration, and reject behavior during active vs idle download states. --- tests/gui/ui_blocks/test_weights_dialog.py | 105 ++++++++++++--------- 1 file changed, 60 insertions(+), 45 deletions(-) diff --git a/tests/gui/ui_blocks/test_weights_dialog.py b/tests/gui/ui_blocks/test_weights_dialog.py index c0812823..9abe06ea 100644 --- a/tests/gui/ui_blocks/test_weights_dialog.py +++ b/tests/gui/ui_blocks/test_weights_dialog.py @@ -1,98 +1,113 @@ +from __future__ import annotations + from PySide6.QtCore import Qt -from PySide6.QtWidgets import QFileDialog +from PySide6.QtWidgets import QDialog from dlclivegui.gui.misc.weights_dialog import PoetWeightsDialog -def test_dialog_without_initial_path_disables_use_button( +def test_dialog_shows_download_destination( + qtbot, +) -> None: + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + assert dialog.path_label.full_text == str(dialog._destination) + assert dialog.path_label.text() + assert dialog.path_label.textInteractionFlags() & Qt.TextInteractionFlag.TextSelectableByMouse + assert dialog.status_label.text() == "Ready to download." + assert dialog.progress_bar.value() == 0 + assert dialog.download_button.isEnabled() + assert dialog.close_button.isEnabled() + + +def test_update_progress_updates_progress_and_status( qtbot, ) -> None: dialog = PoetWeightsDialog() qtbot.addWidget(dialog) - assert dialog.path_label.text() == "No weights selected" - assert not dialog.use_button.isEnabled() - assert not dialog.progress_bar.isVisible() + dialog._update_progress(42) + assert dialog.progress_bar.value() == 42 + assert dialog.status_label.text() == "Downloading POET weights... 42%" -def test_dialog_accepts_valid_initial_path( + +def test_download_finished_records_completed_path( qtbot, tmp_path, ) -> None: path = tmp_path / "weights.pth" path.touch() - dialog = PoetWeightsDialog(initial_path=str(path)) + dialog = PoetWeightsDialog() qtbot.addWidget(dialog) - assert dialog.path_label.text() == str(path) - assert dialog.use_button.isEnabled() + dialog._download_finished(str(path)) + + assert dialog._completed_path == str(path) + assert dialog.progress_bar.value() == 100 + assert dialog.status_label.text() == ("Download complete. Returning to the main window...") -def test_use_selected_emits_path( +def test_download_cleanup_emits_path_and_accepts( qtbot, tmp_path, ) -> None: path = tmp_path / "weights.pth" path.touch() - dialog = PoetWeightsDialog(initial_path=str(path)) + dialog = PoetWeightsDialog() qtbot.addWidget(dialog) + dialog._completed_path = str(path) with qtbot.waitSignal( - dialog.weights_selected, + dialog.weights_downloaded, timeout=1000, ) as blocker: - qtbot.mouseClick( - dialog.use_button, - Qt.MouseButton.LeftButton, - ) + dialog._download_cleanup() assert blocker.args == [str(path)] + assert dialog.result() == QDialog.DialogCode.Accepted + assert dialog._download_thread is None + assert dialog._download_worker is None -def test_browse_sets_selected_path( +def test_download_cleanup_after_failure_restores_controls( qtbot, - monkeypatch, - tmp_path, ) -> None: - path = tmp_path / "weights.pth" - path.touch() + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) - monkeypatch.setattr( - QFileDialog, - "getOpenFileName", - lambda *args, **kwargs: ( - str(path), - "POET weights (*.pt *.pth)", - ), - ) + dialog.download_button.setEnabled(False) + dialog.close_button.setEnabled(False) + dialog._completed_path = None + + dialog._download_cleanup() + + assert dialog.download_button.isEnabled() + assert dialog.close_button.isEnabled() + +def test_reject_is_ignored_while_download_is_active( + qtbot, +) -> None: dialog = PoetWeightsDialog() qtbot.addWidget(dialog) - dialog._browse() + dialog._download_thread = object() + + dialog.reject() - assert dialog.path_label.text() == str(path) - assert dialog.use_button.isEnabled() + assert dialog.result() == 0 -def test_download_cleanup_restores_controls( +def test_reject_closes_when_idle( qtbot, - tmp_path, ) -> None: - path = tmp_path / "weights.pth" - path.touch() - dialog = PoetWeightsDialog() qtbot.addWidget(dialog) - dialog._set_selected_path(str(path)) - dialog.browse_button.setEnabled(False) - dialog.download_button.setEnabled(False) - - dialog._download_cleanup() + dialog.reject() - assert dialog.browse_button.isEnabled() - assert dialog.download_button.isEnabled() - assert dialog.use_button.isEnabled() + assert dialog.result() == QDialog.DialogCode.Rejected From bfa7baa71b45b63cfa2e6495ebf0a3da1abcf8ab Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 15:52:48 +0200 Subject: [PATCH 16/23] Use backend name in inference status messages Replace hardcoded "DLCLive" text in the main window inference startup UI with `_backend_display_name()`. This makes the running button label and initialization success/failure status messages reflect the active backend consistently. --- dlclivegui/gui/main_window.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index eddbcf6f..3b09b955 100644 --- a/dlclivegui/gui/main_window.py +++ b/dlclivegui/gui/main_window.py @@ -3159,16 +3159,16 @@ def _on_pose_processor_initialised(self, success: bool) -> None: self._notify_processor_recording_started(self._processor_recording_context) # Update button to show running state - self.start_inference_button.setText("DLCLive running!") + self.start_inference_button.setText(f"{self._backend_display_name()} running!") self.start_inference_button.setStyleSheet("background-color: #4CAF50; color: white;") - self.statusBar().showMessage("DLCLive initialized successfully", 3000) + self.statusBar().showMessage(f"{self._backend_display_name()} initialized successfully", 3000) return self._dlc_initialized = False # Reset button on failure self.start_inference_button.setText("Start pose inference") self.start_inference_button.setStyleSheet("") - self.statusBar().showMessage("DLCLive initialization failed", 5000) + self.statusBar().showMessage(f"{self._backend_display_name()} initialization failed", 5000) # Stop inference since initialization failed self._stop_inference(show_message=False) From f7aabba6e185e27f2cd49f187513aeaefcfc1d9d Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 16:05:36 +0200 Subject: [PATCH 17/23] Store POET weights in user cache directory Update `poet_default_weights_dir()` to use `~/.cache/dlclivegui/poet` instead of a working-directory temp path. This makes the weights location stable across runs and avoids writing model files into the project directory. --- dlclivegui/services/inference/models/poet/weights.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/dlclivegui/services/inference/models/poet/weights.py b/dlclivegui/services/inference/models/poet/weights.py index f1b2b90a..0718a0bc 100644 --- a/dlclivegui/services/inference/models/poet/weights.py +++ b/dlclivegui/services/inference/models/poet/weights.py @@ -15,9 +15,7 @@ def poet_default_weights_dir() -> Path: - # TODO: move to an app cache rather than a temporary directory - # return Path(tempfile.gettempdir()) / "dlclivegui" / "poet" - return Path.cwd() / "TEMP_WEIGHTS/poet" + return Path.home() / ".cache/dlclivegui/poet" class WeightsDownloadWorker(QObject): From e858214cfad7a45c1a1d44142dd183d3e578b973 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 16:11:21 +0200 Subject: [PATCH 18/23] Lazy-load PyTorch in POET backend Refactors the POET processor to import and validate PyTorch only when POET inference is used, avoiding hard import failures at module load time. It adds explicit runtime errors for missing/incomplete torch installs, moves normalization tensor setup into initialization, replaces the decorator with an inference-mode context, and resets normalization state on shutdown for safer lifecycle handling. --- .../inference/models/poet/poet_processor.py | 90 +++++++++++++++---- 1 file changed, 75 insertions(+), 15 deletions(-) diff --git a/dlclivegui/services/inference/models/poet/poet_processor.py b/dlclivegui/services/inference/models/poet/poet_processor.py index 7cd56b00..25d32b4f 100644 --- a/dlclivegui/services/inference/models/poet/poet_processor.py +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -2,12 +2,13 @@ from __future__ import annotations +import importlib import logging from pathlib import Path +from types import ModuleType from typing import Any import numpy as np -import torch from dlclivegui.services.inference.base import ( PoseBackend, @@ -16,11 +17,45 @@ PoseSource, ) -from .skeleton import POET_KEYPOINT_NAMES, POET_SKELETON_EDGES, POET_SKELETON_ID +from .skeleton import ( + POET_KEYPOINT_NAMES, + POET_SKELETON_EDGES, + POET_SKELETON_ID, +) logger = logging.getLogger(__name__) +def _require_torch() -> ModuleType: + """Import PyTorch when POET inference is actually requested.""" + try: + torch = importlib.import_module("torch") + except ImportError as exc: + raise RuntimeError( + "POET requires PyTorch. Install the POET optional dependencies before using this backend." + ) from exc + + required_attributes = ( + "Tensor", + "autocast", + "cuda", + "device", + "from_numpy", + "inference_mode", + "load", + "tensor", + ) + missing = [name for name in required_attributes if not hasattr(torch, name)] + + if missing: + raise RuntimeError( + "The installed 'torch' module is not a complete PyTorch " + f"installation. Missing attributes: {', '.join(missing)}." + ) + + return torch + + class POETBackend(PoseBackend): """POET pose-inference backend using COCO-17 keypoints.""" @@ -47,15 +82,15 @@ def __init__( self._model: Any | None = None self._postprocessor: Any | None = None - self._device: torch.device | None = None - - self._mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) - self._std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + self._device: Any | None = None + self._mean: Any | None = None + self._std: Any | None = None @staticmethod def _resolve_torch_device( device: str | None, - ) -> torch.device: + ) -> Any: + torch = _require_torch() requested = device.strip().lower() if device else "auto" if requested in {"auto", "best"}: @@ -80,6 +115,11 @@ def init_inference( self, init_frame: np.ndarray, ) -> None: + torch = _require_torch() + + self._mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + self._std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + ( self._model, self._postprocessor, @@ -87,26 +127,30 @@ def init_inference( ) = self._build_model() # Warm up without emitting a result. PoseProcessor emits the - # initial frame's result after init_inference() returns. + # initial frame result after init_inference() returns. self.get_pose(init_frame) def _build_model( self, - ) -> tuple[Any, Any, torch.device]: + ) -> tuple[Any, Any, Any]: + torch = _require_torch() + try: from poet_live import POET, PostProcess - from poet_live.models.backbone import Backbone, Joiner + from poet_live.models.backbone import ( + Backbone, + Joiner, + ) from poet_live.models.position_encoding import ( PositionEmbeddingSine, ) from poet_live.models.transformer import Transformer except ImportError as exc: raise RuntimeError( - "POET is not installed. Install the POET dependencies before using this pose backend." + "POET is not installed. Install the POET optional dependencies before using this backend." ) from exc device = self._resolve_torch_device(self._requested_device) - hidden_dimension = 256 backbone = Backbone( @@ -151,12 +195,14 @@ def _build_model( model.to(device).eval() postprocessor = PostProcess().to(device) + if self._mean is None or self._std is None: + raise RuntimeError("POET normalization tensors are not initialized.") + self._mean = self._mean.to(device) self._std = self._std.to(device) return model, postprocessor, device - @torch.no_grad() def get_pose( self, frame: np.ndarray, @@ -164,18 +210,30 @@ def get_pose( ) -> np.ndarray | None: del frame_time + torch = _require_torch() + + with torch.inference_mode(): + return self._infer(frame, torch) + + def _infer( + self, + frame: np.ndarray, + torch: ModuleType, + ) -> np.ndarray | None: model = self._model postprocessor = self._postprocessor device = self._device + mean = self._mean + std = self._std - if model is None or postprocessor is None or device is None: + if model is None or postprocessor is None or device is None or mean is None or std is None: raise RuntimeError("POET backend is not initialized.") rgb = frame[..., ::-1].copy() height, width = rgb.shape[:2] image = torch.from_numpy(rgb).to(device).permute(2, 0, 1).float().unsqueeze(0) / 255.0 - image = (image - self._mean) / self._std + image = (image - mean) / std use_amp = self._use_amp and device.type == "cuda" @@ -243,3 +301,5 @@ def close(self) -> None: self._model = None self._postprocessor = None self._device = None + self._mean = None + self._std = None From e385c2a20459e92f20f0ba01b87f2209647e4609 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 16:28:03 +0200 Subject: [PATCH 19/23] Handle cancellable POET weights downloads Adds explicit cancellation support for POET weight downloads end-to-end. The weights dialog now requests thread interruption when the window is closed during an active download, shows a cancelling status, and prevents immediate close until the worker finishes. The download worker now checks interruption requests inside the read loop, aborts with a cancellation error, and reliably removes partial `.part` files on failure, while also tightening request/open handling with a timeout and clearer variable naming. --- dlclivegui/gui/misc/weights_dialog.py | 14 +++++ .../services/inference/models/poet/weights.py | 51 +++++++++++++------ 2 files changed, 50 insertions(+), 15 deletions(-) diff --git a/dlclivegui/gui/misc/weights_dialog.py b/dlclivegui/gui/misc/weights_dialog.py index 52f8afcb..4f5c65d2 100644 --- a/dlclivegui/gui/misc/weights_dialog.py +++ b/dlclivegui/gui/misc/weights_dialog.py @@ -163,8 +163,22 @@ def _download_cleanup(self) -> None: self.download_button.setEnabled(True) self.close_button.setEnabled(True) + def request_stop(self) -> None: + thread = self._download_thread + if thread is not None: + thread.requestInterruption() + def reject(self) -> None: if self._download_thread is not None: return super().reject() + + def closeEvent(self, event) -> None: + if self._download_thread is not None: + self.request_stop() + self.status_label.setText("Cancelling download...") + event.ignore() + return + + super().closeEvent(event) diff --git a/dlclivegui/services/inference/models/poet/weights.py b/dlclivegui/services/inference/models/poet/weights.py index 0718a0bc..6281893e 100644 --- a/dlclivegui/services/inference/models/poet/weights.py +++ b/dlclivegui/services/inference/models/poet/weights.py @@ -5,7 +5,7 @@ import urllib.request from pathlib import Path -from PySide6.QtCore import QObject, Signal +from PySide6.QtCore import QObject, QThread, Signal logger = logging.getLogger(__name__) @@ -29,37 +29,58 @@ def __init__(self, url: str, dest: Path): self.dest = dest def run(self) -> None: + tmp = None + try: - self.dest.parent.mkdir(parents=True, exist_ok=True) + self.dest.parent.mkdir( + parents=True, + exist_ok=True, + ) + if self.dest.is_file(): self.progress.emit(100) self.finished.emit(str(self.dest)) return - tmp = self.dest.with_suffix(self.dest.suffix + ".part") - req = urllib.request.Request(self.url, headers={"User-Agent": "DLCLiveGUI"}) - with urllib.request.urlopen(req) as resp, open(tmp, "wb") as f: - total = resp.length or 0 - done = 0 - chunk = 1024 * 256 + tmp = self.dest.with_suffix(self.dest.suffix + ".part") + request = urllib.request.Request( + self.url, + headers={"User-Agent": "DLCLiveGUI"}, + ) + with ( + urllib.request.urlopen( + request, + timeout=30, + ) as response, + tmp.open("wb") as output, + ): + total = response.length or 0 + downloaded = 0 + chunk_size = 256 * 1024 while True: - buf = resp.read(chunk) - if not buf: + if QThread.currentThread().isInterruptionRequested(): + raise RuntimeError("POET weights download was cancelled.") + + chunk = response.read(chunk_size) + if not chunk: break - f.write(buf) - done += len(buf) + + output.write(chunk) + downloaded += len(chunk) + if total > 0: - self.progress.emit(int(done * 100 / total)) + self.progress.emit(int(downloaded * 100 / total)) os.replace(tmp, self.dest) self.progress.emit(100) self.finished.emit(str(self.dest)) - except Exception as e: + except Exception as exc: if tmp is not None: try: tmp.unlink(missing_ok=True) except OSError: logger.exception("Failed to remove partial POET weights.") - self.error.emit(str(e)) + + self.error.emit(str(exc)) From 76d2364448f0cce889bf121d80e1f5f620620e6b Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 28 Aug 2026 10:28:49 +0200 Subject: [PATCH 20/23] Fix queue wait timing and harden POET tests Updated pose processing to stop the loop on queue timeout and compute `queue_wait_time` from enqueue timestamp (`enq`) instead of local get-call timing, improving latency accounting. Expanded POET tests by introducing a reusable `urlopen` stub fixture, validating request/timeout behavior, asserting no `finished` signal on download failure, checking error-free successful downloads, and covering `skeleton_id` metadata in empty-detection pose packets. --- dlclivegui/services/inference/processor.py | 17 +- tests/services/inference/models/test_poet.py | 172 +++++++++++++++---- 2 files changed, 147 insertions(+), 42 deletions(-) diff --git a/dlclivegui/services/inference/processor.py b/dlclivegui/services/inference/processor.py index 93b1dc98..7d1ee348 100644 --- a/dlclivegui/services/inference/processor.py +++ b/dlclivegui/services/inference/processor.py @@ -266,14 +266,21 @@ def _worker_loop(self, init_frame: np.ndarray, init_timestamp: float) -> None: continue try: - wait_start = time.perf_counter() - frame, ts, enq = self._queue.get(timeout=0.05) - qwait = time.perf_counter() - wait_start + frame, ts, enq = self._queue.get( + timeout=0.05, + ) except queue.Empty: - continue + break + else: + queue_wait_time = max(0.0, time.perf_counter() - enq) try: - self._process_frame(frame, ts, enq, queue_wait_time=qwait) + self._process_frame( + frame, + ts, + enq, + queue_wait_time=queue_wait_time, + ) except Exception as exc: logger.exception("Pose inference failed", exc_info=exc) self.error.emit(str(exc)) diff --git a/tests/services/inference/models/test_poet.py b/tests/services/inference/models/test_poet.py index a065d711..b564a4df 100644 --- a/tests/services/inference/models/test_poet.py +++ b/tests/services/inference/models/test_poet.py @@ -1,5 +1,10 @@ +from __future__ import annotations + import urllib.request +from collections.abc import Callable +from dataclasses import dataclass, field from io import BytesIO +from typing import Any import numpy as np import pytest @@ -11,7 +16,93 @@ POET_SKELETON_ID, POETBackend, ) -from dlclivegui.services.inference.models.poet.weights import WeightsDownloadWorker +from dlclivegui.services.inference.models.poet.weights import ( + WeightsDownloadWorker, +) + +TEST_WEIGHTS_URL = "https://example.invalid/weights" + + +class FakeResponse(BytesIO): + """In-memory response implementing the context-manager protocol.""" + + def __init__(self, content: bytes) -> None: + super().__init__(content) + self.length = len(content) + + def __enter__(self) -> FakeResponse: + return self + + def __exit__( + self, + exc_type, + exc_value, + traceback, + ) -> None: + self.close() + + +@dataclass +class UrlopenStub: + """Configurable replacement for urllib.request.urlopen.""" + + content: bytes = b"" + error: Exception | None = None + before_error: Callable[[], None] | None = None + expected_timeout: float = 30 + calls: list[tuple[Any, float | None]] = field(default_factory=list) + + def __call__( + self, + request, + *, + timeout: float | None = None, + ) -> FakeResponse: + self.calls.append((request, timeout)) + + assert timeout == self.expected_timeout + + if self.before_error is not None: + self.before_error() + + if self.error is not None: + raise self.error + + return FakeResponse(self.content) + + +@pytest.fixture +def urlopen_factory( + monkeypatch, +) -> Callable[..., UrlopenStub]: + """Install and return a configurable urlopen test stub.""" + + def install( + *, + content: bytes = b"", + error: Exception | None = None, + before_error: Callable[[], None] | None = None, + expected_timeout: float = 30, + ) -> UrlopenStub: + if error is None and before_error is not None: + raise ValueError("before_error requires an error.") + + stub = UrlopenStub( + content=content, + error=error, + before_error=before_error, + expected_timeout=expected_timeout, + ) + + monkeypatch.setattr( + urllib.request, + "urlopen", + stub, + ) + + return stub + + return install def test_poet_backend_rejects_missing_checkpoint( @@ -44,7 +135,10 @@ def test_make_pose_packet_contains_poet_metadata( checkpoint.touch() backend = POETBackend(str(checkpoint)) - pose = np.zeros((2, 17, 3), dtype=np.float32) + pose = np.zeros( + (2, 17, 3), + dtype=np.float32, + ) packet = backend.make_pose_packet(pose) @@ -72,6 +166,7 @@ def test_make_pose_packet_without_detections_keeps_metadata( assert packet.keypoints is None assert packet.individual_ids == [] assert packet.keypoint_names == list(POET_KEYPOINT_NAMES) + assert packet.skeleton_id == POET_SKELETON_ID assert packet.skeleton_edges == POET_SKELETON_EDGES @@ -83,12 +178,11 @@ def test_download_worker_reuses_existing_file( destination.write_bytes(b"existing") worker = WeightsDownloadWorker( - "https://example.invalid/weights", + TEST_WEIGHTS_URL, destination, ) progress: list[int] = [] - worker.progress.connect(progress.append) with qtbot.waitSignal( @@ -102,76 +196,69 @@ def test_download_worker_reuses_existing_file( assert destination.read_bytes() == b"existing" -class FakeResponse(BytesIO): - def __init__(self, content: bytes) -> None: - super().__init__(content) - self.length = len(content) - - def __enter__(self): - return self - - def __exit__( - self, - exc_type, - exc_value, - traceback, - ) -> None: - self.close() - - def test_download_worker_writes_destination( qtbot, - monkeypatch, tmp_path, + urlopen_factory, ) -> None: destination = tmp_path / "weights.pth" content = b"model-data" - monkeypatch.setattr( - urllib.request, - "urlopen", - lambda _request: FakeResponse(content), + urlopen_stub = urlopen_factory( + content=content, ) worker = WeightsDownloadWorker( - "https://example.invalid/weights", + TEST_WEIGHTS_URL, destination, ) + errors: list[str] = [] + worker.error.connect(errors.append) + with qtbot.waitSignal( worker.finished, timeout=1000, ) as blocker: worker.run() + assert errors == [] assert blocker.args == [str(destination)] assert destination.read_bytes() == content assert not destination.with_suffix(".pth.part").exists() + assert len(urlopen_stub.calls) == 1 + request, timeout = urlopen_stub.calls[0] + + assert isinstance( + request, + urllib.request.Request, + ) + assert request.full_url == TEST_WEIGHTS_URL + assert timeout == 30 + def test_download_worker_reports_failure_and_cleans_partial_file( qtbot, - monkeypatch, tmp_path, + urlopen_factory, ) -> None: destination = tmp_path / "weights.pth" partial_path = destination.with_suffix(".pth.part") - def failing_urlopen(_request): - partial_path.write_bytes(b"partial") - raise OSError("download failed") - - monkeypatch.setattr( - urllib.request, - "urlopen", - failing_urlopen, + urlopen_stub = urlopen_factory( + error=OSError("download failed"), + before_error=lambda: partial_path.write_bytes(b"partial"), ) worker = WeightsDownloadWorker( - "https://example.invalid/weights", + TEST_WEIGHTS_URL, destination, ) + finished_paths: list[str] = [] + worker.finished.connect(finished_paths.append) + with qtbot.waitSignal( worker.error, timeout=1000, @@ -179,5 +266,16 @@ def failing_urlopen(_request): worker.run() assert blocker.args == ["download failed"] + assert finished_paths == [] assert not partial_path.exists() assert not destination.exists() + + assert len(urlopen_stub.calls) == 1 + request, timeout = urlopen_stub.calls[0] + + assert isinstance( + request, + urllib.request.Request, + ) + assert request.full_url == TEST_WEIGHTS_URL + assert timeout == 30 From ed9bc1d9b5e57adb86a0bfb7262516517849ddfa Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 28 Aug 2026 14:58:57 +0200 Subject: [PATCH 21/23] Handle async pose stop and validate POET frames Makes pose processor shutdown non-throwing by returning a success flag from `reset()`/`shutdown()`, and wires that into the main window so users get a status message when stop is deferred during model initialization. It also refactors the worker loop to honor stop requests around backend init, cleanly close backend state in a single `finally` path, and avoid emitting init errors when a stop was intentional. In POET inference, it now validates input as 3-channel image data and enforces contiguous frame memory before inference to prevent backend/runtime issues. --- dlclivegui/gui/main_window.py | 10 +- .../inference/models/poet/poet_processor.py | 7 +- dlclivegui/services/inference/processor.py | 133 ++++++++++-------- 3 files changed, 87 insertions(+), 63 deletions(-) diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index 3b09b955..40af9697 100644 --- a/dlclivegui/gui/main_window.py +++ b/dlclivegui/gui/main_window.py @@ -2984,14 +2984,14 @@ def _reset_active_pose_processor( self, *, reset_processor_plugin: bool, - ) -> None: + ) -> bool: if self._backend_name == "dlc": self._dlc.reset( reset_processor_plugin=reset_processor_plugin, ) - return + return True - self._active_pose_processor.reset() + return self._active_pose_processor.reset() def _stop_inference(self, show_message: bool = True) -> None: if self._rec_manager.is_active and self._backend_name == "dlc": @@ -3011,7 +3011,9 @@ def _stop_inference(self, show_message: bool = True) -> None: self._dlc_active = False self._dlc_initialized = False # Does NOT invoke the normal rec-stop/save hooks. Persistence is processor-dependent. - self._reset_active_pose_processor(reset_processor_plugin=True) + stopped = self._reset_active_pose_processor(reset_processor_plugin=True) + if not stopped: + self.statusBar().showMessage("Stopping pose inference after model init completed...", 5000) self._last_pose = None self._overlay_renderer.clear_runtime_state() self._last_processor_vid_recording = False diff --git a/dlclivegui/services/inference/models/poet/poet_processor.py b/dlclivegui/services/inference/models/poet/poet_processor.py index 25d32b4f..c2aa3306 100644 --- a/dlclivegui/services/inference/models/poet/poet_processor.py +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -210,10 +210,15 @@ def get_pose( ) -> np.ndarray | None: del frame_time + frame = np.asarray(frame) + if frame.ndim != 3 or frame.shape[2] != 3: + # FIXME @C-Achard 2026/08/28 - optionally convert here as well + raise ValueError(f"Incorrect frame format, expected 3-channel BGR, got {frame.shape!r}") + torch = _require_torch() with torch.inference_mode(): - return self._infer(frame, torch) + return self._infer(np.ascontiguousarray(frame), torch) def _infer( self, diff --git a/dlclivegui/services/inference/processor.py b/dlclivegui/services/inference/processor.py index 7d1ee348..7e45f8d9 100644 --- a/dlclivegui/services/inference/processor.py +++ b/dlclivegui/services/inference/processor.py @@ -56,10 +56,12 @@ def configure(self, backend_factory: Callable[[], PoseBackend]) -> None: def is_configured(self) -> bool: return self._backend_factory is not None - def reset(self) -> None: + def reset(self) -> bool: if not self._stop_worker(): - raise RuntimeError("Failed to stop worker thread") + return False + self._backend = None + with self._stats_lock: self._frames_enqueued = 0 self._frames_processed = 0 @@ -71,8 +73,10 @@ def reset(self) -> None: self._signal_emit_times.clear() self._total_process_times.clear() - def shutdown(self) -> None: - self.reset() + return True + + def shutdown(self) -> bool: + return self.reset() def enqueue_frame(self, frame: np.ndarray, timestamp: float) -> None: if self._worker_thread is None: @@ -214,83 +218,96 @@ def _process_frame( self.frame_processed.emit() - def _worker_loop(self, init_frame: np.ndarray, init_timestamp: float) -> None: + def _worker_loop( + self, + init_frame: np.ndarray, + init_timestamp: float, + ) -> None: try: if self._backend_factory is None: raise RuntimeError("No backend configured.") + self._backend = self._backend_factory() + if self._stop_event.is_set(): + return + self._backend.init_inference(init_frame) + + if self._stop_event.is_set(): + return + self.initialized.emit(True) - self._process_frame(init_frame, init_timestamp, time.perf_counter(), queue_wait_time=0.0) + self._process_frame( + init_frame, + init_timestamp, + time.perf_counter(), + queue_wait_time=0.0, + ) + with self._stats_lock: self._frames_enqueued += 1 + while not self._stop_event.is_set(): + queue_ref = self._queue + if queue_ref is None: + break + + try: + frame, timestamp, enqueued_at = queue_ref.get(timeout=0.05) + except queue.Empty: + continue + + queue_wait_time = max( + 0.0, + time.perf_counter() - enqueued_at, + ) + + try: + self._process_frame( + frame, + timestamp, + enqueued_at, + queue_wait_time=queue_wait_time, + ) + except Exception as exc: + logger.exception( + "Pose inference failed", + exc_info=exc, + ) + self.error.emit(str(exc)) + finally: + try: + queue_ref.task_done() + except ValueError: + pass + except Exception as exc: - logger.exception("Failed to initialize pose backend", exc_info=exc) - self.error.emit(str(exc)) - self.initialized.emit(False) + if not self._stop_event.is_set(): + logger.exception( + "Failed to initialize pose backend", + exc_info=exc, + ) + self.error.emit(str(exc)) + self.initialized.emit(False) + finally: backend = self._backend self._backend = None if backend is not None: try: backend.close() - except Exception as exc: - logger.exception("Failed to close backend", exc_info=exc) + except Exception: + logger.exception("Failed to close pose backend.") self._queue = None - self._worker_thread = None - return - while True: - if self._stop_event.is_set(): - if self._queue is not None: - try: - frame, ts, enq = self._queue.get_nowait() - except queue.Empty: - break - else: - try: - self._process_frame(frame, ts, enq, queue_wait_time=0.0) - except Exception as exc: - logger.exception("Pose inference failed", exc_info=exc) - self.error.emit(str(exc)) - finally: - try: - self._queue.task_done() - except ValueError: - pass - continue - - try: - frame, ts, enq = self._queue.get( - timeout=0.05, - ) - except queue.Empty: - break - else: - queue_wait_time = max(0.0, time.perf_counter() - enq) - - try: - self._process_frame( - frame, - ts, - enq, - queue_wait_time=queue_wait_time, - ) - except Exception as exc: - logger.exception("Pose inference failed", exc_info=exc) - self.error.emit(str(exc)) - finally: - try: - self._queue.task_done() - except ValueError: - pass + if self._worker_thread is threading.current_thread(): + self._worker_thread = None - logger.info("Pose worker thread exiting") + logger.info("Pose worker thread exiting") def create_pose_processor( From 74d72046f3f606533d0d187db0f392a83075d834 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 28 Aug 2026 15:29:21 +0200 Subject: [PATCH 22/23] Update reset failure processor tests Revise processor reset tests to match the new non-throwing failure path: reset now returns `False` when the worker does not stop instead of raising. Also add coverage to ensure runtime statistics remain unchanged when reset fails. --- tests/services/inference/test_processor.py | 37 ++++++++++++++++++---- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/tests/services/inference/test_processor.py b/tests/services/inference/test_processor.py index fcaa7dea..5ae44ffc 100644 --- a/tests/services/inference/test_processor.py +++ b/tests/services/inference/test_processor.py @@ -212,7 +212,7 @@ def test_reset_clears_statistics() -> None: assert stats.processing_fps == 0.0 -def test_reset_raises_when_worker_does_not_stop( +def test_reset_returns_false_when_worker_does_not_stop( monkeypatch, ) -> None: processor = PoseProcessor() @@ -223,11 +223,36 @@ def test_reset_raises_when_worker_does_not_stop( lambda: False, ) - with pytest.raises( - RuntimeError, - match="Failed to stop worker thread", - ): - processor.reset() + stopped = processor.reset() + + assert stopped is False + + +def test_failed_reset_preserves_runtime_statistics( + monkeypatch, +) -> None: + processor = PoseProcessor() + processor._frames_enqueued = 4 + processor._frames_processed = 3 + processor._frames_dropped = 2 + processor._latencies.append(1.0) + + monkeypatch.setattr( + processor, + "_stop_worker", + lambda: False, + ) + + stopped = processor.reset() + + assert stopped is False + + stats = processor.get_stats() + + assert stats.frames_enqueued == 4 + assert stats.frames_processed == 3 + assert stats.frames_dropped == 2 + assert stats.average_latency == 1.0 def test_create_dlc_pose_processor() -> None: From 63e490ac65e714d515a9bf826b3bc63f1acb15d7 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Fri, 28 Aug 2026 15:33:15 +0200 Subject: [PATCH 23/23] Harden POET checkpoint loading Update POET model loading to use `torch.serialization.safe_globals([argparse.Namespace])` and `weights_only=True` when calling `torch.load`. This reduces unsafe pickle deserialization risk while still allowing legacy checkpoint metadata needed by the model state load. --- .../inference/models/poet/poet_processor.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/dlclivegui/services/inference/models/poet/poet_processor.py b/dlclivegui/services/inference/models/poet/poet_processor.py index c2aa3306..e631c74c 100644 --- a/dlclivegui/services/inference/models/poet/poet_processor.py +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -2,6 +2,7 @@ from __future__ import annotations +import argparse import importlib import logging from pathlib import Path @@ -182,11 +183,13 @@ def _build_model( aux_loss=False, ) - checkpoint = torch.load( - self._checkpoint_path, - map_location="cpu", - weights_only=False, - ) + with torch.serialization.safe_globals([argparse.Namespace]): + checkpoint = torch.load( + self._checkpoint_path, + map_location="cpu", + # weights_only=False, + weights_only=True, + ) model.load_state_dict( checkpoint["model"], strict=True,