From de034ab9e63308c79fc8ad5f268ec9bad195f8bd Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 08:44:35 +0200 Subject: [PATCH 1/8] Create skeleton.py --- dlclivegui/utils/skeleton.py | 367 +++++++++++++++++++++++++++++++++++ 1 file changed, 367 insertions(+) create mode 100644 dlclivegui/utils/skeleton.py diff --git a/dlclivegui/utils/skeleton.py b/dlclivegui/utils/skeleton.py new file mode 100644 index 000000000..1bacbf28b --- /dev/null +++ b/dlclivegui/utils/skeleton.py @@ -0,0 +1,367 @@ +"""Skeleton definition, validation, and drawing utilities.""" + +# dlclivegui/utils/skeleton.py +from __future__ import annotations + +import json +from dataclasses import dataclass +from enum import Enum, auto +from pathlib import Path + +import cv2 +import numpy as np +import yaml +from pydantic import BaseModel, Field, ValidationError, field_validator + +from dlclivegui.config import BGR, SkeletonColorMode, SkeletonStyle + + +# ############### # +# Status & code # +# ############### # +class SkeletonRenderCode(Enum): + OK = auto() + POSE_SHAPE_INVALID = auto() + KEYPOINT_COUNT_MISMATCH = auto() + + +@dataclass(frozen=True) +class SkeletonRenderStatus: + code: SkeletonRenderCode + message: str = "" + + @property + def rendered(self) -> bool: + return self.code == SkeletonRenderCode.OK + + @property + def should_disable(self) -> bool: + # GUI can switch off skeleton drawing if True + return self.code in { + SkeletonRenderCode.POSE_SHAPE_INVALID, + SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, + } + + +# ############ # +# Exceptions # +# ############ # + + +class SkeletonError(ValueError): + """Raised when a skeleton definition is invalid.""" + + +class SkeletonLoadError(Exception): + """High-level skeleton loading error (safe for GUI display).""" + + +class SkeletonValidationError(SkeletonLoadError): + """Schema or semantic validation error.""" + + +# ################## # +# Skeleton display # +# ################## # + + +class SkeletonStyleModel(BaseModel): + mode: SkeletonColorMode = SkeletonColorMode.SOLID + color: BGR = (0, 255, 255) # default if SOLID + thickness: int = Field(2, ge=1, description="Base thickness in pixels") + gradient_steps: int = Field(16, ge=2, description="Segments per edge when gradient") + scale_with_zoom: bool = True + + @field_validator("thickness") + @classmethod + def _thickness_positive(cls, v): + if v < 1: + raise ValueError("Thickness must be at least 1 pixel") + return v + + @field_validator("gradient_steps") + @classmethod + def _steps_positive(cls, v): + if v < 2: + raise ValueError("gradient_steps must be >= 2") + return v + + +# ############# # +# Skeleton IO # +# ############# # +class SkeletonModel(BaseModel): + """Validated skeleton definition (IO + schema).""" + + name: str | None = None + + keypoints: list[str] = Field(..., min_length=1, description="Ordered list of keypoint names") + + edges: list[tuple[int, int]] = Field( + default_factory=list, + description="List of (i, j) keypoint index pairs", + ) + + style: SkeletonStyleModel = Field(default_factory=SkeletonStyleModel) + default_color: BGR = (0, 255, 255) # used if style.color is None or in SOLID mode + edge_colors: dict[tuple[int, int], BGR] = Field(default_factory=dict) + + schema_version: int = 1 + + @field_validator("keypoints") + @classmethod + def validate_unique_keypoints(cls, v): + if len(set(v)) != len(v): + raise ValueError("Duplicate keypoint names detected") + return v + + @field_validator("edges") + @classmethod + def validate_edges(cls, edges, info): + keypoints = info.data.get("keypoints", []) + n = len(keypoints) + + for i, j in edges: + if i == j: + raise ValueError(f"Self-loop detected in edge ({i}, {j})") + if not (0 <= i < n and 0 <= j < n): + raise ValueError(f"Edge ({i}, {j}) out of range for {n} keypoints") + return edges + + +def _load_raw_skeleton_data(path: Path) -> dict: + if not path.exists(): + raise SkeletonLoadError(f"Skeleton file not found: {path}") + + if path.suffix in {".yaml", ".yml"}: + return yaml.safe_load(path.read_text()) + + if path.suffix == ".json": + return json.loads(path.read_text()) + + raise SkeletonLoadError(f"Unsupported file type: {path.suffix}") + + +def _format_pydantic_error(err: ValidationError) -> str: + lines = ["Invalid skeleton definition:"] + for e in err.errors(): + loc = " → ".join(map(str, e["loc"])) + msg = e["msg"] + lines.append(f"• {loc}: {msg}") + return "\n".join(lines) + + +def load_skeleton(path: Path) -> Skeleton: + try: + data = _load_raw_skeleton_data(path) + model = SkeletonModel.model_validate(data) + return Skeleton(model) + + except ValidationError as e: + raise SkeletonValidationError(_format_pydantic_error(e)) from None + + except Exception as e: + raise SkeletonLoadError(str(e)) from None + + +def save_skeleton(path: Path, model: SkeletonModel) -> None: + data = model.model_dump() + + if path.suffix in {".yaml", ".yml"}: + path.write_text(yaml.safe_dump(data, sort_keys=False)) + elif path.suffix == ".json": + path.write_text(json.dumps(data, indent=2)) + else: + raise SkeletonLoadError(f"Unsupported skeleton file type: {path.suffix}") + + +def load_dlc_skeleton(config_path: Path) -> Skeleton | None: + if not config_path.exists(): + raise SkeletonLoadError(f"DLC config not found: {config_path}") + + cfg = yaml.safe_load(config_path.read_text()) + + bodyparts = cfg.get("bodyparts") + if not bodyparts: + return None # No pose info + + edges = [] + + # Newer DLC format + if "skeleton" in cfg: + for a, b in cfg["skeleton"]: + edges.append((bodyparts.index(a), bodyparts.index(b))) + + # Older / alternative formats + elif "skeleton_edges" in cfg: + edges = [tuple(e) for e in cfg["skeleton_edges"]] + + if not edges: + return None + + model = SkeletonModel( + name=cfg.get("Task", "DeepLabCut"), + keypoints=bodyparts, + edges=edges, + ) + + return Skeleton(model) + + +class Skeleton: + """Runtime skeleton optimized for drawing.""" + + def __init__(self, model: SkeletonModel): + self.name = model.name + self.keypoints = model.keypoints + self.edges = model.edges + + self.style = SkeletonStyle( + mode=model.style.mode, + color=model.style.color, + thickness=model.style.thickness, + gradient_steps=model.style.gradient_steps, + scale_with_zoom=model.style.scale_with_zoom, + ) + self.default_color = model.default_color + self.edge_colors = model.edge_colors + + def check_pose_compat(self, pose: np.ndarray) -> SkeletonRenderStatus: + pose = np.asarray(pose) + + if pose.ndim != 2 or pose.shape[1] not in (2, 3): + return SkeletonRenderStatus( + SkeletonRenderCode.POSE_SHAPE_INVALID, + f"Pose must be (N,2) or (N,3); got shape={pose.shape}", + ) + + expected = len(self.keypoints) + got = pose.shape[0] + if got != expected: + return SkeletonRenderStatus( + SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, + f"Skeleton expects {expected} keypoints, but pose has {got}.", + ) + + return SkeletonRenderStatus(SkeletonRenderCode.OK, "") + + def _draw_gradient_edge( + self, + img: np.ndarray, + p1: tuple[int, int], + p2: tuple[int, int], + c1: BGR, + c2: BGR, + thickness: int, + steps: int, + ): + x1, y1 = p1 + x2, y2 = p2 + + for s in range(steps): + a0 = s / steps + a1 = (s + 1) / steps + xs0 = int(x1 + (x2 - x1) * a0) + ys0 = int(y1 + (y2 - y1) * a0) + xs1 = int(x1 + (x2 - x1) * a1) + ys1 = int(y1 + (y2 - y1) * a1) + + t = (s + 0.5) / steps + b = int(c1[0] + (c2[0] - c1[0]) * t) + g = int(c1[1] + (c2[1] - c1[1]) * t) + r = int(c1[2] + (c2[2] - c1[2]) * t) + + cv2.line(img, (xs0, ys0), (xs1, ys1), (b, g, r), thickness, lineType=cv2.LINE_AA) + + def draw( + self, + overlay: np.ndarray, + pose: np.ndarray, + p_cutoff: float, + offset: tuple[int, int], + scale: tuple[float, float], + *, + style: SkeletonStyle | None = None, + color_override: BGR | None = None, + keypoint_colors: list[BGR] | None = None, + ) -> SkeletonRenderStatus: + status = self.check_pose_compat(pose) + if not status.rendered: + return status + + st = style or self.style + ox, oy = offset + sx, sy = scale + th = st.effective_thickness(sx, sy) + + # if gradient mode, require keypoint_colors aligned with keypoint order + if st.mode == SkeletonColorMode.GRADIENT_KEYPOINTS: + if keypoint_colors is None or len(keypoint_colors) != len(self.keypoints): + return SkeletonRenderStatus( + SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, + f"Gradient mode requires keypoint_colors of length {len(self.keypoints)}.", + ) + + for i, j in self.edges: + xi, yi = pose[i][:2] + xj, yj = pose[j][:2] + ci = pose[i][2] if pose.shape[1] > 2 else 1.0 + cj = pose[j][2] if pose.shape[1] > 2 else 1.0 + if np.isnan(xi) or np.isnan(yi) or ci < p_cutoff or np.isnan(xj) or np.isnan(yj) or cj < p_cutoff: + continue + + p1 = (int(xi * sx + ox), int(yi * sy + oy)) + p2 = (int(xj * sx + ox), int(yj * sy + oy)) + + if st.mode == SkeletonColorMode.GRADIENT_KEYPOINTS: + c1 = keypoint_colors[i] + c2 = keypoint_colors[j] + self._draw_gradient_edge(overlay, p1, p2, c1, c2, th, st.gradient_steps) + else: + # SOLID: priority edge_colors > override > style.color > default_color + color = self.edge_colors.get((i, j), color_override or st.color or self.default_color) + cv2.line(overlay, p1, p2, color, th, lineType=cv2.LINE_AA) + + return SkeletonRenderStatus(SkeletonRenderCode.OK, "") + + def draw_many( + self, + overlay: np.ndarray, + poses: np.ndarray, + p_cutoff: float, + offset: tuple[int, int], + scale: tuple[float, float], + *, + style: SkeletonStyle | None = None, + color_override: BGR | None = None, + keypoint_colors: list[BGR] | None = None, + ) -> SkeletonRenderStatus: + poses = np.asarray(poses) + if poses.ndim != 3: + return SkeletonRenderStatus( + SkeletonRenderCode.POSE_SHAPE_INVALID, + f"Multi-pose must be (A,N,2/3); got shape={poses.shape}", + ) + + expected = len(self.keypoints) + if poses.shape[1] != expected: + return SkeletonRenderStatus( + SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, + f"Skeleton expects {expected} keypoints, but poses have N={poses.shape[1]}.", + ) + + for pose in poses: + st = self.draw( + overlay, + pose, + p_cutoff, + offset, + scale, + style=style, + color_override=color_override, + keypoint_colors=keypoint_colors, + ) + if not st.rendered: + return st + + return SkeletonRenderStatus(SkeletonRenderCode.OK, "") From 71ad37a384c59aa39480b23ab825d938535ab7d0 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 08:47:11 +0200 Subject: [PATCH 2/8] Add skeleton config fields --- dlclivegui/config.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/dlclivegui/config.py b/dlclivegui/config.py index 11bcf8049..905e78ebb 100644 --- a/dlclivegui/config.py +++ b/dlclivegui/config.py @@ -12,6 +12,7 @@ from dlclivegui.utils.writegear_options import WriteGearOptions Rotation = Literal[0, 90, 180, 270] +BGR = tuple[int, int, int] # color format TileLayout = Literal["auto", "2x2", "1x4", "4x1"] Precision = Literal["FP32", "FP16"] ModelType = Literal["pytorch", "tensorflow"] @@ -500,6 +501,24 @@ def _bbox_logic(self): return self +class SkeletonColorMode(str, Enum): + SOLID = "solid" + GRADIENT_KEYPOINTS = "gradient_keypoints" # use endpoint keypoint colors + + +class SkeletonStyle(BaseModel): + mode: SkeletonColorMode = SkeletonColorMode.SOLID + color: BGR = (0, 255, 255) # default if SOLID + thickness: int = 2 # base thickness in pixels + gradient_steps: int = 16 # segments per edge when gradient + scale_with_zoom: bool = True # scale thickness with (sx, sy) + + def effective_thickness(self, sx: float, sy: float) -> int: + if not self.scale_with_zoom: + return max(1, int(self.thickness)) + return max(1, int(round(self.thickness * min(sx, sy)))) + + class VisualizationSettings(BaseModel): p_cutoff: float = Field(default=0.6, ge=0.0, le=1.0) colormap: str = "hot" From d1d7f7fa7c38d2ba728c20617a51d6e1c35ed108 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 10:39:24 +0200 Subject: [PATCH 3/8] Move display and skeleton code --- dlclivegui/display/__init__.py | 19 +++++++++++++++++++ dlclivegui/{utils => display}/display.py | 0 dlclivegui/{utils => display}/skeleton.py | 0 dlclivegui/gui/main_window.py | 2 +- tests/utils/test_display.py | 2 +- 5 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 dlclivegui/display/__init__.py rename dlclivegui/{utils => display}/display.py (100%) rename dlclivegui/{utils => display}/skeleton.py (100%) diff --git a/dlclivegui/display/__init__.py b/dlclivegui/display/__init__.py new file mode 100644 index 000000000..e7e64b5df --- /dev/null +++ b/dlclivegui/display/__init__.py @@ -0,0 +1,19 @@ +from .display import ( + BBoxColors, + compute_tile_info, + compute_tiling_geometry, + create_tiled_frame, + draw_bbox, + draw_keypoints, + draw_pose, +) + +__all__ = [ + "BBoxColors", + "compute_tile_info", + "compute_tiling_geometry", + "create_tiled_frame", + "draw_bbox", + "draw_keypoints", + "draw_pose", +] diff --git a/dlclivegui/utils/display.py b/dlclivegui/display/display.py similarity index 100% rename from dlclivegui/utils/display.py rename to dlclivegui/display/display.py diff --git a/dlclivegui/utils/skeleton.py b/dlclivegui/display/skeleton.py similarity index 100% rename from dlclivegui/utils/skeleton.py rename to dlclivegui/display/skeleton.py diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index abb877021..6e056b74a 100644 --- a/dlclivegui/gui/main_window.py +++ b/dlclivegui/gui/main_window.py @@ -65,6 +65,7 @@ VisualizationSettings, ) +from ..display import BBoxColors, compute_tile_info, create_tiled_frame, draw_bbox, draw_pose from ..processors.processor_utils import ( create_spec_from_scan, default_processors_dir, @@ -77,7 +78,6 @@ from ..services.dlc_processor import DLCLiveProcessor, PoseResult from ..services.multi_camera_controller import MultiCameraController, MultiFrameData, get_camera_id, get_display_id from ..services.recording_manager import RecordingManager -from ..utils.display import BBoxColors, compute_tile_info, create_tiled_frame, draw_bbox, draw_pose from ..utils.settings_store import DLCLiveGUISettingsStore, ModelPathStore from ..utils.stats import WorkerTimingStats, format_dlc_stats from ..utils.utils import FPSTracker diff --git a/tests/utils/test_display.py b/tests/utils/test_display.py index 559aa1522..ec6b7eec8 100644 --- a/tests/utils/test_display.py +++ b/tests/utils/test_display.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from dlclivegui.utils.display import ( # noqa: E402 +from dlclivegui.display import ( # noqa: E402 compute_tile_info, compute_tiling_geometry, create_tiled_frame, From ec96fe73bbffe1c60548b1047247cd7a7bc4eca5 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 10:49:51 +0200 Subject: [PATCH 4/8] Refine SkeletonStyle config fields Update `SkeletonStyle` to support explicit visibility and clearer naming (`color_mode`, `color_bgr`), and switch numeric styling options to constrained `Field` definitions for validated thickness and gradient step ranges. --- dlclivegui/config.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/dlclivegui/config.py b/dlclivegui/config.py index 905e78ebb..e8f5c3492 100644 --- a/dlclivegui/config.py +++ b/dlclivegui/config.py @@ -507,10 +507,11 @@ class SkeletonColorMode(str, Enum): class SkeletonStyle(BaseModel): - mode: SkeletonColorMode = SkeletonColorMode.SOLID - color: BGR = (0, 255, 255) # default if SOLID - thickness: int = 2 # base thickness in pixels - gradient_steps: int = 16 # segments per edge when gradient + visible: bool = False + color_mode: SkeletonColorMode = SkeletonColorMode.SOLID + color_bgr: BGR = (0, 255, 255) # default if SOLID + thickness: int = Field(defalt=2, ge=1, le=20) # base thickness in pixels + gradient_steps: int = Field(default=16, ge=2, le=128) # segments per edge when gradient scale_with_zoom: bool = True # scale thickness with (sx, sy) def effective_thickness(self, sx: float, sy: float) -> int: From 3a4a4936c78dbcde892c33857033496e5ce49e17 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 10:50:06 +0200 Subject: [PATCH 5/8] Add skeleton resolver and style field updates Introduce typed skeleton resolution primitives (`SkeletonEdge`, `SkeletonDefinition`, `ResolvedSkeleton`) and a `resolve_skeleton` helper that validates keypoint inputs, detects duplicates/missing names, and raises `SkeletonResolutionError` with clear messages. Also align rendering with renamed style fields (`mode` -> `color_mode`, `color` -> `color_bgr`) so initialization and draw paths use the updated `SkeletonStyle` API, and fix the module path header comment. --- dlclivegui/display/skeleton.py | 76 +++++++++++++++++++++++++++++++--- 1 file changed, 70 insertions(+), 6 deletions(-) diff --git a/dlclivegui/display/skeleton.py b/dlclivegui/display/skeleton.py index 1bacbf28b..04967590e 100644 --- a/dlclivegui/display/skeleton.py +++ b/dlclivegui/display/skeleton.py @@ -1,6 +1,6 @@ """Skeleton definition, validation, and drawing utilities.""" -# dlclivegui/utils/skeleton.py +# dlclivegui/display/skeleton.py from __future__ import annotations import json @@ -43,11 +43,75 @@ def should_disable(self) -> bool: } +@dataclass(frozen=True, slots=True) +class SkeletonEdge: + start: str + end: str + + +@dataclass(frozen=True, slots=True) +class SkeletonDefinition: + identifier: str + display_name: str + edges: tuple[SkeletonEdge, ...] + + +@dataclass(frozen=True, slots=True) +class ResolvedSkeleton: + definition: SkeletonDefinition + keypoint_names: tuple[str, ...] + edges: tuple[tuple[int, int], ...] + + +def resolve_skeleton( + definition: SkeletonDefinition, + keypoint_names: list[str] | tuple[str, ...], +) -> ResolvedSkeleton: + names = tuple(keypoint_names) + + if not names: + raise SkeletonResolutionError("Cannot resolve a skeleton without keypoint names.") + + if len(set(names)) != len(names): + raise SkeletonResolutionError("Cannot resolve a skeleton against duplicate keypoint names.") + + name_to_index = {name: index for index, name in enumerate(names)} + + resolved_edges: list[tuple[int, int]] = [] + missing_names: set[str] = set() + + for edge in definition.edges: + start_index = name_to_index.get(edge.start) + end_index = name_to_index.get(edge.end) + + if start_index is None: + missing_names.add(edge.start) + if end_index is None: + missing_names.add(edge.end) + + if start_index is not None and end_index is not None: + resolved_edges.append((start_index, end_index)) + + if missing_names: + missing = ", ".join(sorted(missing_names)) + raise SkeletonResolutionError(f"Skeleton keypoints are absent from the pose output: {missing}.") + + return ResolvedSkeleton( + definition=definition, + keypoint_names=names, + edges=tuple(resolved_edges), + ) + + # ############ # # Exceptions # # ############ # +class SkeletonResolutionError(ValueError): + """Raised when a skeleton cannot be aligned with pose keypoints.""" + + class SkeletonError(ValueError): """Raised when a skeleton definition is invalid.""" @@ -217,8 +281,8 @@ def __init__(self, model: SkeletonModel): self.edges = model.edges self.style = SkeletonStyle( - mode=model.style.mode, - color=model.style.color, + color_mode=model.style.mode, + color_bgr=model.style.color, thickness=model.style.thickness, gradient_steps=model.style.gradient_steps, scale_with_zoom=model.style.scale_with_zoom, @@ -295,7 +359,7 @@ def draw( th = st.effective_thickness(sx, sy) # if gradient mode, require keypoint_colors aligned with keypoint order - if st.mode == SkeletonColorMode.GRADIENT_KEYPOINTS: + if st.color_mode == SkeletonColorMode.GRADIENT_KEYPOINTS: if keypoint_colors is None or len(keypoint_colors) != len(self.keypoints): return SkeletonRenderStatus( SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, @@ -313,13 +377,13 @@ def draw( p1 = (int(xi * sx + ox), int(yi * sy + oy)) p2 = (int(xj * sx + ox), int(yj * sy + oy)) - if st.mode == SkeletonColorMode.GRADIENT_KEYPOINTS: + if st.color_mode == SkeletonColorMode.GRADIENT_KEYPOINTS: c1 = keypoint_colors[i] c2 = keypoint_colors[j] self._draw_gradient_edge(overlay, p1, p2, c1, c2, th, st.gradient_steps) else: # SOLID: priority edge_colors > override > style.color > default_color - color = self.edge_colors.get((i, j), color_override or st.color or self.default_color) + color = self.edge_colors.get((i, j), color_override or st.color_bgr or self.default_color) cv2.line(overlay, p1, p2, color, th, lineType=cv2.LINE_AA) return SkeletonRenderStatus(SkeletonRenderCode.OK, "") From 358eb13fc5b89a843aef3d3db44f7084996c54d5 Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 11:00:08 +0200 Subject: [PATCH 6/8] Extract shared inference base types Move pose-related enums and dataclasses (PoseBackends, WorkerState, PoseSource, PosePacket, PoseResult) out of `dlc_processor.py` into a new `services/inference/base.py` module and import them back into the processor. This centralizes inference contracts for reuse and extends `PosePacket` with optional skeleton metadata fields (`skeleton_id`, `skeleton_edges`). --- dlclivegui/services/dlc_processor.py | 40 ++------------------- dlclivegui/services/inference/__init__.py | 0 dlclivegui/services/inference/base.py | 44 +++++++++++++++++++++++ 3 files changed, 47 insertions(+), 37 deletions(-) create mode 100644 dlclivegui/services/inference/__init__.py create mode 100644 dlclivegui/services/inference/base.py diff --git a/dlclivegui/services/dlc_processor.py b/dlclivegui/services/dlc_processor.py index c2993d791..6002fdd4f 100644 --- a/dlclivegui/services/dlc_processor.py +++ b/dlclivegui/services/dlc_processor.py @@ -10,13 +10,12 @@ from collections import deque from contextlib import contextmanager from dataclasses import dataclass -from enum import Enum, auto from typing import Any import numpy as np from PySide6.QtCore import QObject, Signal -from dlclivegui.config import DLC_DO_LOG_TIMING, DLCProcessorSettings, ModelType +from dlclivegui.config import DLC_DO_LOG_TIMING, DLCProcessorSettings from dlclivegui.processors.processor_utils import ( ProcessorSpec, create_spec_from_scan, @@ -28,6 +27,8 @@ from dlclivegui.utils.stats import WorkerTimingStats from dlclivegui.utils.utils import format_thread_stack +from .inference.base import PoseBackends, PosePacket, PoseResult, PoseSource, WorkerState + logger = logging.getLogger(__name__) STOP_WORKER_TIMEOUT = 10.0 # # seconds to wait in STOPPING state before scheduling background reaping @@ -44,41 +45,6 @@ ENABLE_PROFILING = False -class PoseBackends(Enum): - DLC_LIVE = auto() - - -class WorkerState(Enum): - STOPPED = auto() - STARTING = auto() - RUNNING = auto() - STOPPING = auto() - FAULTED = auto() - - -@dataclass -class PoseResult: - pose: np.ndarray | None - timestamp: float - packet: PosePacket | None = None - - -@dataclass(slots=True, frozen=True) -class PoseSource: - backend: PoseBackends # e.g. "DLCLive" - model_type: ModelType | None = None - - -@dataclass(slots=True, frozen=True) -class PosePacket: - schema_version: int = 0 - keypoints: np.ndarray | None = None - keypoint_names: list[str] | None = None - individual_ids: list[str] | None = None - source: PoseSource = PoseSource(backend=PoseBackends.DLC_LIVE) - raw: Any | None = None - - def validate_pose_array( pose: Any, *, source_backend: PoseBackends | str = PoseBackends.DLC_LIVE, check_finite: bool = True ) -> np.ndarray: diff --git a/dlclivegui/services/inference/__init__.py b/dlclivegui/services/inference/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/dlclivegui/services/inference/base.py b/dlclivegui/services/inference/base.py new file mode 100644 index 000000000..35a5d9d56 --- /dev/null +++ b/dlclivegui/services/inference/base.py @@ -0,0 +1,44 @@ +from dataclasses import dataclass +from enum import Enum, auto +from typing import Any + +import numpy as np + +from dlclivegui.config import ModelType + + +class PoseBackends(Enum): + DLC_LIVE = auto() + + +class WorkerState(Enum): + STOPPED = auto() + STARTING = auto() + RUNNING = auto() + STOPPING = auto() + FAULTED = auto() + + +@dataclass(slots=True, frozen=True) +class PoseSource: + backend: PoseBackends # e.g. "DLCLive" + model_type: ModelType | None = None + + +@dataclass(slots=True, frozen=True) +class PosePacket: + schema_version: int = 0 + keypoints: np.ndarray | None = None + keypoint_names: list[str] | None = None + individual_ids: list[str] | None = None + skeleton_id: str | None = None + skeleton_edges: tuple[tuple[str, str], ...] | None = None + source: PoseSource = PoseSource(backend=PoseBackends.DLC_LIVE) + raw: Any | None = None + + +@dataclass +class PoseResult: + pose: np.ndarray | None + timestamp: float + packet: PosePacket | None = None From 1b22395d90ae3a2adea30d4e5314a2f10ee80d3a Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 11:00:41 +0200 Subject: [PATCH 7/8] Refactor skeleton rendering API Reworks `display/skeleton.py`. It replaces the old `Skeleton` runtime/model loading stack with immutable skeleton definition/resolution dataclasses, adds validated metadata-to-definition construction, and introduces a standalone `draw_skeleton(...)` path that supports both single and multi-individual poses. Rendering now returns a structured `SkeletonRenderResult` with clearer status codes (including no-pose, invalid shape, keypoint mismatch, and gradient color mismatch) and tracks drawn edge count. --- dlclivegui/display/skeleton.py | 537 +++++++++++++-------------------- 1 file changed, 218 insertions(+), 319 deletions(-) diff --git a/dlclivegui/display/skeleton.py b/dlclivegui/display/skeleton.py index 04967590e..80c9fc4c3 100644 --- a/dlclivegui/display/skeleton.py +++ b/dlclivegui/display/skeleton.py @@ -1,56 +1,36 @@ -"""Skeleton definition, validation, and drawing utilities.""" +"""Skeleton topology resolution and rendering utilities.""" -# dlclivegui/display/skeleton.py from __future__ import annotations -import json from dataclasses import dataclass from enum import Enum, auto -from pathlib import Path import cv2 import numpy as np -import yaml -from pydantic import BaseModel, Field, ValidationError, field_validator from dlclivegui.config import BGR, SkeletonColorMode, SkeletonStyle +# ####################### # +# Skeleton definitions # +# ####################### # -# ############### # -# Status & code # -# ############### # -class SkeletonRenderCode(Enum): - OK = auto() - POSE_SHAPE_INVALID = auto() - KEYPOINT_COUNT_MISMATCH = auto() - - -@dataclass(frozen=True) -class SkeletonRenderStatus: - code: SkeletonRenderCode - message: str = "" - @property - def rendered(self) -> bool: - return self.code == SkeletonRenderCode.OK - - @property - def should_disable(self) -> bool: - # GUI can switch off skeleton drawing if True - return self.code in { - SkeletonRenderCode.POSE_SHAPE_INVALID, - SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, - } +class SkeletonResolutionError(ValueError): + """Raised when a skeleton cannot be aligned with pose keypoints.""" @dataclass(frozen=True, slots=True) class SkeletonEdge: + """An edge expressed using semantic keypoint names.""" + start: str end: str @dataclass(frozen=True, slots=True) class SkeletonDefinition: + """Immutable, backend-independent skeleton topology.""" + identifier: str display_name: str edges: tuple[SkeletonEdge, ...] @@ -58,6 +38,8 @@ class SkeletonDefinition: @dataclass(frozen=True, slots=True) class ResolvedSkeleton: + """Skeleton topology resolved against a specific keypoint order.""" + definition: SkeletonDefinition keypoint_names: tuple[str, ...] edges: tuple[tuple[int, int], ...] @@ -67,6 +49,7 @@ def resolve_skeleton( definition: SkeletonDefinition, keypoint_names: list[str] | tuple[str, ...], ) -> ResolvedSkeleton: + """Resolve named skeleton edges against an ordered keypoint list.""" names = tuple(keypoint_names) if not names: @@ -76,7 +59,6 @@ def resolve_skeleton( raise SkeletonResolutionError("Cannot resolve a skeleton against duplicate keypoint names.") name_to_index = {name: index for index, name in enumerate(names)} - resolved_edges: list[tuple[int, int]] = [] missing_names: set[str] = set() @@ -103,329 +85,246 @@ def resolve_skeleton( ) -# ############ # -# Exceptions # -# ############ # - - -class SkeletonResolutionError(ValueError): - """Raised when a skeleton cannot be aligned with pose keypoints.""" - - -class SkeletonError(ValueError): - """Raised when a skeleton definition is invalid.""" - - -class SkeletonLoadError(Exception): - """High-level skeleton loading error (safe for GUI display).""" - +# ###################### # +# Skeleton I/O # +# ###################### # -class SkeletonValidationError(SkeletonLoadError): - """Schema or semantic validation error.""" +def skeleton_definition_from_metadata( + *, + identifier: str, + display_name: str, + edges: tuple[tuple[str, str], ...], +) -> SkeletonDefinition: + """Create a validated display topology from pose metadata.""" + if not identifier.strip(): + raise SkeletonResolutionError("Skeleton identifier cannot be empty.") -# ################## # -# Skeleton display # -# ################## # - - -class SkeletonStyleModel(BaseModel): - mode: SkeletonColorMode = SkeletonColorMode.SOLID - color: BGR = (0, 255, 255) # default if SOLID - thickness: int = Field(2, ge=1, description="Base thickness in pixels") - gradient_steps: int = Field(16, ge=2, description="Segments per edge when gradient") - scale_with_zoom: bool = True - - @field_validator("thickness") - @classmethod - def _thickness_positive(cls, v): - if v < 1: - raise ValueError("Thickness must be at least 1 pixel") - return v - - @field_validator("gradient_steps") - @classmethod - def _steps_positive(cls, v): - if v < 2: - raise ValueError("gradient_steps must be >= 2") - return v + if not edges: + raise SkeletonResolutionError("Skeleton definition does not contain any edges.") + skeleton_edges: list[SkeletonEdge] = [] -# ############# # -# Skeleton IO # -# ############# # -class SkeletonModel(BaseModel): - """Validated skeleton definition (IO + schema).""" + for start, end in edges: + if not start or not end: + raise SkeletonResolutionError("Skeleton edge names cannot be empty.") - name: str | None = None + if start == end: + raise SkeletonResolutionError(f"Skeleton contains a self-loop at {start!r}.") - keypoints: list[str] = Field(..., min_length=1, description="Ordered list of keypoint names") + skeleton_edges.append( + SkeletonEdge( + start=start, + end=end, + ) + ) - edges: list[tuple[int, int]] = Field( - default_factory=list, - description="List of (i, j) keypoint index pairs", + return SkeletonDefinition( + identifier=identifier, + display_name=display_name, + edges=tuple(skeleton_edges), ) - style: SkeletonStyleModel = Field(default_factory=SkeletonStyleModel) - default_color: BGR = (0, 255, 255) # used if style.color is None or in SOLID mode - edge_colors: dict[tuple[int, int], BGR] = Field(default_factory=dict) - - schema_version: int = 1 - - @field_validator("keypoints") - @classmethod - def validate_unique_keypoints(cls, v): - if len(set(v)) != len(v): - raise ValueError("Duplicate keypoint names detected") - return v - - @field_validator("edges") - @classmethod - def validate_edges(cls, edges, info): - keypoints = info.data.get("keypoints", []) - n = len(keypoints) - for i, j in edges: - if i == j: - raise ValueError(f"Self-loop detected in edge ({i}, {j})") - if not (0 <= i < n and 0 <= j < n): - raise ValueError(f"Edge ({i}, {j}) out of range for {n} keypoints") - return edges +# ###################### # +# Rendering outcomes # +# ###################### # -def _load_raw_skeleton_data(path: Path) -> dict: - if not path.exists(): - raise SkeletonLoadError(f"Skeleton file not found: {path}") - - if path.suffix in {".yaml", ".yml"}: - return yaml.safe_load(path.read_text()) - - if path.suffix == ".json": - return json.loads(path.read_text()) - - raise SkeletonLoadError(f"Unsupported file type: {path.suffix}") +class SkeletonRenderCode(Enum): + RENDERED = auto() + NO_POSE = auto() + INVALID_POSE = auto() + KEYPOINT_COUNT_MISMATCH = auto() + COLOR_COUNT_MISMATCH = auto() -def _format_pydantic_error(err: ValidationError) -> str: - lines = ["Invalid skeleton definition:"] - for e in err.errors(): - loc = " → ".join(map(str, e["loc"])) - msg = e["msg"] - lines.append(f"• {loc}: {msg}") - return "\n".join(lines) +@dataclass(frozen=True, slots=True) +class SkeletonRenderResult: + code: SkeletonRenderCode + edges_drawn: int = 0 + message: str = "" + @property + def rendered(self) -> bool: + return self.code == SkeletonRenderCode.RENDERED + + +# ###################### # +# Rendering utilities # +# ###################### # + + +def _effective_thickness( + style: SkeletonStyle, + scale: tuple[float, float], +) -> int: + scale_x, scale_y = scale + return style.effective_thickness(scale_x, scale_y) + + +def _draw_gradient_edge( + frame: np.ndarray, + start: tuple[int, int], + end: tuple[int, int], + start_color: BGR, + end_color: BGR, + *, + thickness: int, + steps: int, +) -> None: + start_x, start_y = start + end_x, end_y = end + + for step in range(steps): + alpha_start = step / steps + alpha_end = (step + 1) / steps + color_alpha = (step + 0.5) / steps + + segment_start = ( + round(start_x + (end_x - start_x) * alpha_start), + round(start_y + (end_y - start_y) * alpha_start), + ) + segment_end = ( + round(start_x + (end_x - start_x) * alpha_end), + round(start_y + (end_y - start_y) * alpha_end), + ) -def load_skeleton(path: Path) -> Skeleton: - try: - data = _load_raw_skeleton_data(path) - model = SkeletonModel.model_validate(data) - return Skeleton(model) + color: BGR = tuple( + round(component_start + (component_end - component_start) * color_alpha) + for component_start, component_end in zip( + start_color, + end_color, + strict=True, + ) + ) - except ValidationError as e: - raise SkeletonValidationError(_format_pydantic_error(e)) from None + cv2.line( + frame, + segment_start, + segment_end, + color, + thickness, + lineType=cv2.LINE_AA, + ) - except Exception as e: - raise SkeletonLoadError(str(e)) from None +def draw_skeleton( + frame: np.ndarray, + poses: np.ndarray | None, + skeleton: ResolvedSkeleton, + style: SkeletonStyle, + *, + p_cutoff: float, + offset: tuple[int, int] = (0, 0), + scale: tuple[float, float] = (1.0, 1.0), + keypoint_colors: tuple[BGR, ...] | None = None, +) -> SkeletonRenderResult: + """Draw a resolved skeleton over one or more poses. + + Accepted pose shapes are: + + - ``(K, 3)`` for one individual + - ``(N, K, 3)`` for multiple individuals + + The function modifies ``frame`` in place and returns a structured result. + """ + if poses is None: + return SkeletonRenderResult( + code=SkeletonRenderCode.NO_POSE, + ) -def save_skeleton(path: Path, model: SkeletonModel) -> None: - data = model.model_dump() + pose_array = np.asarray(poses) - if path.suffix in {".yaml", ".yml"}: - path.write_text(yaml.safe_dump(data, sort_keys=False)) - elif path.suffix == ".json": - path.write_text(json.dumps(data, indent=2)) + if pose_array.ndim == 2: + individuals = pose_array[np.newaxis, ...] + elif pose_array.ndim == 3: + individuals = pose_array else: - raise SkeletonLoadError(f"Unsupported skeleton file type: {path.suffix}") - - -def load_dlc_skeleton(config_path: Path) -> Skeleton | None: - if not config_path.exists(): - raise SkeletonLoadError(f"DLC config not found: {config_path}") - - cfg = yaml.safe_load(config_path.read_text()) - - bodyparts = cfg.get("bodyparts") - if not bodyparts: - return None # No pose info - - edges = [] - - # Newer DLC format - if "skeleton" in cfg: - for a, b in cfg["skeleton"]: - edges.append((bodyparts.index(a), bodyparts.index(b))) - - # Older / alternative formats - elif "skeleton_edges" in cfg: - edges = [tuple(e) for e in cfg["skeleton_edges"]] - - if not edges: - return None - - model = SkeletonModel( - name=cfg.get("Task", "DeepLabCut"), - keypoints=bodyparts, - edges=edges, - ) - - return Skeleton(model) - + return SkeletonRenderResult( + code=SkeletonRenderCode.INVALID_POSE, + message=(f"Skeleton poses must have shape (K, 3) or (N, K, 3); received {pose_array.shape!r}."), + ) -class Skeleton: - """Runtime skeleton optimized for drawing.""" + if individuals.shape[-1] != 3: + return SkeletonRenderResult( + code=SkeletonRenderCode.INVALID_POSE, + message=(f"Skeleton poses must contain x, y, and likelihood; received {pose_array.shape!r}."), + ) - def __init__(self, model: SkeletonModel): - self.name = model.name - self.keypoints = model.keypoints - self.edges = model.edges + expected_keypoints = len(skeleton.keypoint_names) + actual_keypoints = individuals.shape[1] - self.style = SkeletonStyle( - color_mode=model.style.mode, - color_bgr=model.style.color, - thickness=model.style.thickness, - gradient_steps=model.style.gradient_steps, - scale_with_zoom=model.style.scale_with_zoom, + if actual_keypoints != expected_keypoints: + return SkeletonRenderResult( + code=SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, + message=(f"Skeleton expects {expected_keypoints} keypoints, but the pose contains {actual_keypoints}."), ) - self.default_color = model.default_color - self.edge_colors = model.edge_colors - def check_pose_compat(self, pose: np.ndarray) -> SkeletonRenderStatus: - pose = np.asarray(pose) + uses_gradient = style.color_mode == SkeletonColorMode.GRADIENT_KEYPOINTS - if pose.ndim != 2 or pose.shape[1] not in (2, 3): - return SkeletonRenderStatus( - SkeletonRenderCode.POSE_SHAPE_INVALID, - f"Pose must be (N,2) or (N,3); got shape={pose.shape}", - ) + if uses_gradient and (keypoint_colors is None or len(keypoint_colors) != expected_keypoints): + return SkeletonRenderResult( + code=SkeletonRenderCode.COLOR_COUNT_MISMATCH, + message=(f"Keypoint-gradient mode requires exactly {expected_keypoints} keypoint colors."), + ) - expected = len(self.keypoints) - got = pose.shape[0] - if got != expected: - return SkeletonRenderStatus( - SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, - f"Skeleton expects {expected} keypoints, but pose has {got}.", + offset_x, offset_y = offset + scale_x, scale_y = scale + thickness = _effective_thickness(style, scale) + edges_drawn = 0 + + for pose in individuals: + for start_index, end_index in skeleton.edges: + start_x, start_y, start_likelihood = pose[start_index] + end_x, end_y, end_likelihood = pose[end_index] + + values = ( + start_x, + start_y, + start_likelihood, + end_x, + end_y, + end_likelihood, ) - return SkeletonRenderStatus(SkeletonRenderCode.OK, "") - - def _draw_gradient_edge( - self, - img: np.ndarray, - p1: tuple[int, int], - p2: tuple[int, int], - c1: BGR, - c2: BGR, - thickness: int, - steps: int, - ): - x1, y1 = p1 - x2, y2 = p2 - - for s in range(steps): - a0 = s / steps - a1 = (s + 1) / steps - xs0 = int(x1 + (x2 - x1) * a0) - ys0 = int(y1 + (y2 - y1) * a0) - xs1 = int(x1 + (x2 - x1) * a1) - ys1 = int(y1 + (y2 - y1) * a1) - - t = (s + 0.5) / steps - b = int(c1[0] + (c2[0] - c1[0]) * t) - g = int(c1[1] + (c2[1] - c1[1]) * t) - r = int(c1[2] + (c2[2] - c1[2]) * t) - - cv2.line(img, (xs0, ys0), (xs1, ys1), (b, g, r), thickness, lineType=cv2.LINE_AA) - - def draw( - self, - overlay: np.ndarray, - pose: np.ndarray, - p_cutoff: float, - offset: tuple[int, int], - scale: tuple[float, float], - *, - style: SkeletonStyle | None = None, - color_override: BGR | None = None, - keypoint_colors: list[BGR] | None = None, - ) -> SkeletonRenderStatus: - status = self.check_pose_compat(pose) - if not status.rendered: - return status - - st = style or self.style - ox, oy = offset - sx, sy = scale - th = st.effective_thickness(sx, sy) - - # if gradient mode, require keypoint_colors aligned with keypoint order - if st.color_mode == SkeletonColorMode.GRADIENT_KEYPOINTS: - if keypoint_colors is None or len(keypoint_colors) != len(self.keypoints): - return SkeletonRenderStatus( - SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, - f"Gradient mode requires keypoint_colors of length {len(self.keypoints)}.", - ) - - for i, j in self.edges: - xi, yi = pose[i][:2] - xj, yj = pose[j][:2] - ci = pose[i][2] if pose.shape[1] > 2 else 1.0 - cj = pose[j][2] if pose.shape[1] > 2 else 1.0 - if np.isnan(xi) or np.isnan(yi) or ci < p_cutoff or np.isnan(xj) or np.isnan(yj) or cj < p_cutoff: + if not np.isfinite(values).all() or start_likelihood < p_cutoff or end_likelihood < p_cutoff: continue - p1 = (int(xi * sx + ox), int(yi * sy + oy)) - p2 = (int(xj * sx + ox), int(yj * sy + oy)) - - if st.color_mode == SkeletonColorMode.GRADIENT_KEYPOINTS: - c1 = keypoint_colors[i] - c2 = keypoint_colors[j] - self._draw_gradient_edge(overlay, p1, p2, c1, c2, th, st.gradient_steps) - else: - # SOLID: priority edge_colors > override > style.color > default_color - color = self.edge_colors.get((i, j), color_override or st.color_bgr or self.default_color) - cv2.line(overlay, p1, p2, color, th, lineType=cv2.LINE_AA) - - return SkeletonRenderStatus(SkeletonRenderCode.OK, "") - - def draw_many( - self, - overlay: np.ndarray, - poses: np.ndarray, - p_cutoff: float, - offset: tuple[int, int], - scale: tuple[float, float], - *, - style: SkeletonStyle | None = None, - color_override: BGR | None = None, - keypoint_colors: list[BGR] | None = None, - ) -> SkeletonRenderStatus: - poses = np.asarray(poses) - if poses.ndim != 3: - return SkeletonRenderStatus( - SkeletonRenderCode.POSE_SHAPE_INVALID, - f"Multi-pose must be (A,N,2/3); got shape={poses.shape}", + start_point = ( + round(start_x * scale_x + offset_x), + round(start_y * scale_y + offset_y), ) - - expected = len(self.keypoints) - if poses.shape[1] != expected: - return SkeletonRenderStatus( - SkeletonRenderCode.KEYPOINT_COUNT_MISMATCH, - f"Skeleton expects {expected} keypoints, but poses have N={poses.shape[1]}.", + end_point = ( + round(end_x * scale_x + offset_x), + round(end_y * scale_y + offset_y), ) - for pose in poses: - st = self.draw( - overlay, - pose, - p_cutoff, - offset, - scale, - style=style, - color_override=color_override, - keypoint_colors=keypoint_colors, - ) - if not st.rendered: - return st + if uses_gradient: + assert keypoint_colors is not None + + _draw_gradient_edge( + frame, + start_point, + end_point, + keypoint_colors[start_index], + keypoint_colors[end_index], + thickness=thickness, + steps=style.gradient_steps, + ) + else: + cv2.line( + frame, + start_point, + end_point, + style.color_bgr, + thickness, + lineType=cv2.LINE_AA, + ) - return SkeletonRenderStatus(SkeletonRenderCode.OK, "") + edges_drawn += 1 + + return SkeletonRenderResult( + code=SkeletonRenderCode.RENDERED, + edges_drawn=edges_drawn, + ) From f63f14d2596993defc30b91b2edc23fb84336bff Mon Sep 17 00:00:00 2001 From: Cyril Achard Date: Thu, 27 Aug 2026 11:09:41 +0200 Subject: [PATCH 8/8] Refactor shared stats and skeleton parsing Move `ProcessorStats` from `dlc_processor` into `services.inference.base` so it can be imported as a shared inference type, and update dependent imports/tests accordingly. Add a typed `SkeletonPacket` protocol plus `resolve_packet_skeleton()` to consistently resolve packet-provided skeleton metadata into a `ResolvedSkeleton` only when required fields are present. --- dlclivegui/display/skeleton.py | 39 +++++++++++++++++++++++---- dlclivegui/services/dlc_processor.py | 24 +---------------- dlclivegui/services/inference/base.py | 21 +++++++++++++++ dlclivegui/utils/stats.py | 2 +- tests/services/test_dlc_processor.py | 2 ++ 5 files changed, 59 insertions(+), 29 deletions(-) diff --git a/dlclivegui/display/skeleton.py b/dlclivegui/display/skeleton.py index 80c9fc4c3..8da0edcb3 100644 --- a/dlclivegui/display/skeleton.py +++ b/dlclivegui/display/skeleton.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from enum import Enum, auto +from typing import Protocol import cv2 import numpy as np @@ -85,11 +86,6 @@ def resolve_skeleton( ) -# ###################### # -# Skeleton I/O # -# ###################### # - - def skeleton_definition_from_metadata( *, identifier: str, @@ -126,6 +122,39 @@ def skeleton_definition_from_metadata( ) +# ###################### # +# Skeleton I/O # +# ###################### # + + +class SkeletonPacket(Protocol): + keypoint_names: list[str] | None + skeleton_id: str | None + skeleton_edges: tuple[tuple[str, str], ...] | None + + +def resolve_packet_skeleton( + packet: SkeletonPacket, +) -> ResolvedSkeleton | None: + """Resolve skeleton metadata supplied by a pose packet.""" + if not packet.keypoint_names: + return None + + if not packet.skeleton_id or not packet.skeleton_edges: + return None + + definition = skeleton_definition_from_metadata( + identifier=packet.skeleton_id, + display_name=packet.skeleton_id, + edges=packet.skeleton_edges, + ) + + return resolve_skeleton( + definition, + packet.keypoint_names, + ) + + # ###################### # # Rendering outcomes # # ###################### # diff --git a/dlclivegui/services/dlc_processor.py b/dlclivegui/services/dlc_processor.py index 6002fdd4f..09d255256 100644 --- a/dlclivegui/services/dlc_processor.py +++ b/dlclivegui/services/dlc_processor.py @@ -9,7 +9,6 @@ import time from collections import deque from contextlib import contextmanager -from dataclasses import dataclass from typing import Any import numpy as np @@ -27,7 +26,7 @@ from dlclivegui.utils.stats import WorkerTimingStats from dlclivegui.utils.utils import format_thread_stack -from .inference.base import PoseBackends, PosePacket, PoseResult, PoseSource, WorkerState +from .inference.base import PoseBackends, PosePacket, PoseResult, PoseSource, ProcessorStats, WorkerState logger = logging.getLogger(__name__) STOP_WORKER_TIMEOUT = 10.0 # # seconds to wait in STOPPING state before scheduling background reaping @@ -93,27 +92,6 @@ def validate_pose_array( return arr -@dataclass -class ProcessorStats: - """Statistics for DLC processor performance.""" - - frames_enqueued: int = 0 - frames_processed: int = 0 - frames_dropped: int = 0 - queue_size: int = 0 - processing_fps: float = 0.0 - average_latency: float = 0.0 - last_latency: float = 0.0 - # Profiling metrics - avg_queue_wait: float = 0.0 - avg_inference_time: float = 0.0 - avg_signal_emit_time: float = 0.0 - avg_total_process_time: float = 0.0 - # 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 DLCLiveProcessor(QObject): """Background pose estimation using DLCLive with queue-based threading.""" diff --git a/dlclivegui/services/inference/base.py b/dlclivegui/services/inference/base.py index 35a5d9d56..17d0286b8 100644 --- a/dlclivegui/services/inference/base.py +++ b/dlclivegui/services/inference/base.py @@ -42,3 +42,24 @@ class PoseResult: pose: np.ndarray | None timestamp: float packet: PosePacket | None = None + + +@dataclass +class ProcessorStats: + """Statistics for DLC processor performance.""" + + frames_enqueued: int = 0 + frames_processed: int = 0 + frames_dropped: int = 0 + queue_size: int = 0 + processing_fps: float = 0.0 + average_latency: float = 0.0 + last_latency: float = 0.0 + # Profiling metrics + avg_queue_wait: float = 0.0 + avg_inference_time: float = 0.0 + avg_signal_emit_time: float = 0.0 + avg_total_process_time: float = 0.0 + # 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 diff --git a/dlclivegui/utils/stats.py b/dlclivegui/utils/stats.py index 0ef0528e2..48ebe23e6 100644 --- a/dlclivegui/utils/stats.py +++ b/dlclivegui/utils/stats.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from dlclivegui.services.dlc_processor import ProcessorStats + from dlclivegui.services.inference.base import ProcessorStats @dataclass diff --git a/tests/services/test_dlc_processor.py b/tests/services/test_dlc_processor.py index 42b72d718..28b192b19 100644 --- a/tests/services/test_dlc_processor.py +++ b/tests/services/test_dlc_processor.py @@ -12,6 +12,8 @@ from dlclivegui.processors.processor_utils import ProcessorSpec from dlclivegui.services.dlc_processor import ( DLCLiveProcessor, +) +from dlclivegui.services.inference.base import ( ProcessorStats, WorkerState, )