diff --git a/dlclivegui/config.py b/dlclivegui/config.py index 11bcf8049..e8f5c3492 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,25 @@ def _bbox_logic(self): return self +class SkeletonColorMode(str, Enum): + SOLID = "solid" + GRADIENT_KEYPOINTS = "gradient_keypoints" # use endpoint keypoint colors + + +class SkeletonStyle(BaseModel): + 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: + 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" 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/display/skeleton.py b/dlclivegui/display/skeleton.py new file mode 100644 index 000000000..8da0edcb3 --- /dev/null +++ b/dlclivegui/display/skeleton.py @@ -0,0 +1,359 @@ +"""Skeleton topology resolution and rendering utilities.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum, auto +from typing import Protocol + +import cv2 +import numpy as np + +from dlclivegui.config import BGR, SkeletonColorMode, SkeletonStyle + +# ####################### # +# Skeleton definitions # +# ####################### # + + +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, ...] + + +@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], ...] + + +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: + 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), + ) + + +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.") + + if not edges: + raise SkeletonResolutionError("Skeleton definition does not contain any edges.") + + skeleton_edges: list[SkeletonEdge] = [] + + for start, end in edges: + if not start or not end: + raise SkeletonResolutionError("Skeleton edge names cannot be empty.") + + if start == end: + raise SkeletonResolutionError(f"Skeleton contains a self-loop at {start!r}.") + + skeleton_edges.append( + SkeletonEdge( + start=start, + end=end, + ) + ) + + return SkeletonDefinition( + identifier=identifier, + display_name=display_name, + edges=tuple(skeleton_edges), + ) + + +# ###################### # +# 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 # +# ###################### # + + +class SkeletonRenderCode(Enum): + RENDERED = auto() + NO_POSE = auto() + INVALID_POSE = auto() + KEYPOINT_COUNT_MISMATCH = auto() + COLOR_COUNT_MISMATCH = auto() + + +@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), + ) + + 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, + ) + ) + + cv2.line( + frame, + segment_start, + segment_end, + color, + thickness, + lineType=cv2.LINE_AA, + ) + + +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, + ) + + pose_array = np.asarray(poses) + + if pose_array.ndim == 2: + individuals = pose_array[np.newaxis, ...] + elif pose_array.ndim == 3: + individuals = pose_array + else: + return SkeletonRenderResult( + code=SkeletonRenderCode.INVALID_POSE, + message=(f"Skeleton poses must have shape (K, 3) or (N, K, 3); received {pose_array.shape!r}."), + ) + + 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}."), + ) + + expected_keypoints = len(skeleton.keypoint_names) + actual_keypoints = individuals.shape[1] + + 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}."), + ) + + uses_gradient = style.color_mode == SkeletonColorMode.GRADIENT_KEYPOINTS + + 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."), + ) + + 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, + ) + + if not np.isfinite(values).all() or start_likelihood < p_cutoff or end_likelihood < p_cutoff: + continue + + start_point = ( + round(start_x * scale_x + offset_x), + round(start_y * scale_y + offset_y), + ) + end_point = ( + round(end_x * scale_x + offset_x), + round(end_y * scale_y + offset_y), + ) + + 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, + ) + + edges_drawn += 1 + + return SkeletonRenderResult( + code=SkeletonRenderCode.RENDERED, + edges_drawn=edges_drawn, + ) 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/dlclivegui/services/dlc_processor.py b/dlclivegui/services/dlc_processor.py index c2993d791..09d255256 100644 --- a/dlclivegui/services/dlc_processor.py +++ b/dlclivegui/services/dlc_processor.py @@ -9,14 +9,12 @@ import time 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 +26,8 @@ from dlclivegui.utils.stats import WorkerTimingStats from dlclivegui.utils.utils import format_thread_stack +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 @@ -44,41 +44,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: @@ -127,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/__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..17d0286b8 --- /dev/null +++ b/dlclivegui/services/inference/base.py @@ -0,0 +1,65 @@ +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 + + +@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, ) 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,