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/config.py b/dlclivegui/config.py index 14d00250..0558e7cc 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 +MODEL_INFERENCE_PROFILING_ENABLED: bool = False # MAIN_WINDOW_DO_LOG_TIMING: bool = False #### Backends BASLER_DO_LOG_TIMING: bool = False diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py index 6ecb496c..40af9697 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, cast 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 @@ -93,10 +95,13 @@ 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 logger = logging.getLogger("DLCLiveGUI") +if TYPE_CHECKING: + from ..services.inference.models.poet.poet_processor import POETBackend class DLCLiveMainWindow(QMainWindow): @@ -148,7 +153,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( @@ -184,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 @@ -210,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() @@ -372,7 +383,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 +405,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) @@ -420,6 +431,70 @@ 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) + + dialog.weights_downloaded.connect(self._on_poet_weights_downloaded) + 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_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, + ) -> 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") @@ -470,6 +545,42 @@ 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) + # ------------------------- + 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) + + 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 +608,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 +707,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.""" @@ -810,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) @@ -850,36 +988,49 @@ 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) + return group - 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) + # ------------------------------------------------------------------ signals + def _connect_pose_processor_signals(self) -> None: + processor = self._active_pose_processor - 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) + 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, + ) - 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) + def _disconnect_pose_processor_signals(self) -> None: + processor = self._active_pose_processor - form.addRow("Coordinates", bbox_layout) + connections = ( + ( + processor.pose_ready, + self._on_pose_ready, + ), + ( + processor.error, + self._on_pose_error, + ), + ( + processor.initialized, + self._on_pose_processor_initialised, + ), + ) - return group + for signal, slot in connections: + try: + signal.disconnect(slot) + except RuntimeError: + pass - # ------------------------------------------------------------------ signals def _connect_signals(self) -> None: self.preview_button.clicked.connect(self._start_preview) self.stop_preview_button.clicked.connect(self._stop_preview) @@ -915,9 +1066,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) @@ -941,7 +1090,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)) @@ -1050,6 +1202,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 +1228,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 +1455,72 @@ 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) + 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())) @@ -1410,7 +1585,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 +2058,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 +2146,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 +2206,140 @@ 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 _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: + + 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: + self._show_warning("Stop pose inference before switching backends.") + self._sync_backend_actions() + return + + self._remember_current_backend_path() + self._disconnect_pose_processor_signals() + + self._backend_name = backend + 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._restore_current_backend_path() + + self.settings.setValue( + "app/backend", + backend, + ) + + self.statusBar().showMessage( + f"Pose backend set to {self._backend_display_name()}", + 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" + + 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 +2364,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 +2382,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 +2457,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 +2694,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 +2770,15 @@ 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) + self.action_manage_poet_weights.setEnabled(allow_changes) def _update_camera_controls_enabled(self) -> None: multi_cam_recording = self._rec_manager.is_active @@ -2492,11 +2861,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 +2882,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 +2955,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, + ) -> bool: + if self._backend_name == "dlc": + self._dlc.reset( + reset_processor_plugin=reset_processor_plugin, + ) + return True + + 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 +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._dlc.reset(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 @@ -2704,7 +3098,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 +3150,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.setText(f"{self._backend_display_name()} 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) + 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(f"{self._backend_display_name()} initialization failed", 5000) + # Stop inference since initialization failed + self._stop_inference(show_message=False) # ------------------------------------------------------------------ # Helpers @@ -2817,6 +3216,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 +3252,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() @@ -2851,7 +3260,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/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 new file mode 100644 index 00000000..4f5c65d2 --- /dev/null +++ b/dlclivegui/gui/misc/weights_dialog.py @@ -0,0 +1,184 @@ +"""Dialog for downloading the default POET weights.""" + +from __future__ import annotations + +from PySide6.QtCore import QThread, Signal, Slot +from PySide6.QtWidgets import ( + QDialog, + QHBoxLayout, + QLabel, + QMessageBox, + QProgressBar, + QPushButton, + QVBoxLayout, +) + +from dlclivegui.gui.misc.eliding_label import ElidingPathLabel +from dlclivegui.services.inference.models.poet.weights import ( + POET_WEIGHTS_FILENAME, + POET_WEIGHTS_URL, + WeightsDownloadWorker, + poet_default_weights_dir, +) + + +class PoetWeightsDialog(QDialog): + """Download the default POET checkpoint and return its path.""" + + weights_downloaded = Signal(str) + + def __init__(self, parent=None) -> None: + super().__init__(parent) + self.setWindowTitle("Download POET weights") + self.setModal(False) + self.setMinimumWidth(480) + + 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) + + 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.setTextVisible(True) + + self.download_button = QPushButton("Download weights") + self.close_button = QPushButton("Close") + + button_row = QHBoxLayout() + 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) + + self.download_button.clicked.connect(self._download) + self.close_button.clicked.connect(self.reject) + + @Slot() + def _download(self) -> None: + if self._download_thread is not None: + return + + thread = QThread(self) + worker = WeightsDownloadWorker( + POET_WEIGHTS_URL, + self._destination, + ) + worker.moveToThread(thread) + + thread.started.connect(worker.run) + + worker.progress.connect(self._update_progress) + 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._completed_path = None + + self.status_label.setText("Downloading POET weights...") + self.progress_bar.setValue(0) + + self.download_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.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", + message, + ) + + @Slot() + def _download_cleanup(self) -> None: + self._download_thread = None + self._download_worker = None + + completed_path = self._completed_path + self._completed_path = None + + if completed_path is not None: + self.weights_downloaded.emit(completed_path) + self.accept() + return + + 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/base.py b/dlclivegui/services/inference/base.py index 17d0286b..891e5792 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,9 +8,12 @@ from dlclivegui.config import ModelType +logger = logging.getLogger(__name__) -class PoseBackends(Enum): - DLC_LIVE = auto() + +class PoseBackends(str, Enum): + DLC_LIVE = "DLC_LIVE" + POET = "POET" class WorkerState(Enum): @@ -63,3 +68,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.""" 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..e631c74c --- /dev/null +++ b/dlclivegui/services/inference/models/poet/poet_processor.py @@ -0,0 +1,313 @@ +"""POET pose-estimation backend.""" + +from __future__ import annotations + +import argparse +import importlib +import logging +from pathlib import Path +from types import ModuleType +from typing import Any + +import numpy as np + +from dlclivegui.services.inference.base import ( + PoseBackend, + PoseBackends, + PosePacket, + PoseSource, +) + +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.""" + + 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: Any | None = None + self._mean: Any | None = None + self._std: Any | None = None + + @staticmethod + def _resolve_torch_device( + device: str | None, + ) -> Any: + torch = _require_torch() + 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: + 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, + self._device, + ) = self._build_model() + + # Warm up without emitting a result. PoseProcessor emits the + # initial frame result after init_inference() returns. + self.get_pose(init_frame) + + def _build_model( + self, + ) -> 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.position_encoding import ( + PositionEmbeddingSine, + ) + from poet_live.models.transformer import Transformer + except ImportError as exc: + raise RuntimeError( + "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( + "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, + ) + + 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, + ) + + 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 + + def get_pose( + self, + frame: np.ndarray, + frame_time: float | None = None, + ) -> 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(np.ascontiguousarray(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 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 - mean) / 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_SKELETON_ID, + 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 + self._mean = None + self._std = 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..e3284589 --- /dev/null +++ b/dlclivegui/services/inference/models/poet/skeleton.py @@ -0,0 +1,48 @@ +"""POET COCO-17 keypoint and skeleton metadata.""" + +from __future__ import annotations + +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_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 new file mode 100644 index 00000000..6281893e --- /dev/null +++ b/dlclivegui/services/inference/models/poet/weights.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import logging +import os +import urllib.request +from pathlib import Path + +from PySide6.QtCore import QObject, QThread, Signal + +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: + return Path.home() / ".cache/dlclivegui/poet" + + +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: + tmp = 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") + 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: + if QThread.currentThread().isInterruptionRequested(): + raise RuntimeError("POET weights download was cancelled.") + + chunk = response.read(chunk_size) + if not chunk: + break + + output.write(chunk) + downloaded += len(chunk) + + if total > 0: + 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 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(exc)) diff --git a/dlclivegui/services/inference/processor.py b/dlclivegui/services/inference/processor.py new file mode 100644 index 00000000..7e45f8d9 --- /dev/null +++ b/dlclivegui/services/inference/processor.py @@ -0,0 +1,323 @@ +import logging +import queue +import threading +import time +from collections import deque +from collections.abc import Callable +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): + """ + 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[[], PoseBackend] | None = None + self._backend: PoseBackend | 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[[], PoseBackend]) -> None: + self._backend_factory = backend_factory + + def is_configured(self) -> bool: + return self._backend_factory is not None + + def reset(self) -> bool: + if not self._stop_worker(): + return False + + 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() + + 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: + 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) -> bool: + worker = self._worker_thread + if worker is None: + 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 False + + 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.") + return True + + 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 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) + 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() + + 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, + ) + + 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: + 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: + logger.exception("Failed to close pose backend.") + + self._queue = None + + if self._worker_thread is threading.current_thread(): + self._worker_thread = None + + 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}.") 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): 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 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..9abe06ea --- /dev/null +++ b/tests/gui/ui_blocks/test_weights_dialog.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from PySide6.QtCore import Qt +from PySide6.QtWidgets import QDialog + +from dlclivegui.gui.misc.weights_dialog import PoetWeightsDialog + + +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) + + dialog._update_progress(42) + + assert dialog.progress_bar.value() == 42 + assert dialog.status_label.text() == "Downloading POET weights... 42%" + + +def test_download_finished_records_completed_path( + qtbot, + tmp_path, +) -> None: + path = tmp_path / "weights.pth" + path.touch() + + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + 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_download_cleanup_emits_path_and_accepts( + qtbot, + tmp_path, +) -> None: + path = tmp_path / "weights.pth" + path.touch() + + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + dialog._completed_path = str(path) + + with qtbot.waitSignal( + dialog.weights_downloaded, + timeout=1000, + ) as blocker: + 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_download_cleanup_after_failure_restores_controls( + qtbot, +) -> None: + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + 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._download_thread = object() + + dialog.reject() + + assert dialog.result() == 0 + + +def test_reject_closes_when_idle( + qtbot, +) -> None: + dialog = PoetWeightsDialog() + qtbot.addWidget(dialog) + + dialog.reject() + + assert dialog.result() == QDialog.DialogCode.Rejected diff --git a/tests/services/inference/models/test_poet.py b/tests/services/inference/models/test_poet.py new file mode 100644 index 00000000..b564a4df --- /dev/null +++ b/tests/services/inference/models/test_poet.py @@ -0,0 +1,281 @@ +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 + +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, +) + +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( + 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_id == POET_SKELETON_ID + 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( + TEST_WEIGHTS_URL, + 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" + + +def test_download_worker_writes_destination( + qtbot, + tmp_path, + urlopen_factory, +) -> None: + destination = tmp_path / "weights.pth" + content = b"model-data" + + urlopen_stub = urlopen_factory( + content=content, + ) + + worker = WeightsDownloadWorker( + 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, + tmp_path, + urlopen_factory, +) -> None: + destination = tmp_path / "weights.pth" + partial_path = destination.with_suffix(".pth.part") + + urlopen_stub = urlopen_factory( + error=OSError("download failed"), + before_error=lambda: partial_path.write_bytes(b"partial"), + ) + + worker = WeightsDownloadWorker( + TEST_WEIGHTS_URL, + destination, + ) + + finished_paths: list[str] = [] + worker.finished.connect(finished_paths.append) + + with qtbot.waitSignal( + worker.error, + timeout=1000, + ) as blocker: + 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 diff --git a/tests/services/inference/test_processor.py b/tests/services/inference/test_processor.py new file mode 100644 index 00000000..5ae44ffc --- /dev/null +++ b/tests/services/inference/test_processor.py @@ -0,0 +1,275 @@ +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_returns_false_when_worker_does_not_stop( + monkeypatch, +) -> None: + processor = PoseProcessor() + + monkeypatch.setattr( + processor, + "_stop_worker", + lambda: False, + ) + + 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: + 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")