Use Engine enum and improve model detection

Introduce a temporary Engine enum and expose it via dlclivegui.temp to standardize model engine types. Update DLCProcessorSettings to use the Engine enum for model_type. Improve model backend detection in the main window by calling DLCLiveProcessor.get_model_backend and raising a clear error if detection fails; tighten file dialog filters to focus on PyTorch models. Enhance engine detection to recognize .pth files. Minor import cleanup in dlc_processor and a clarifying test comment about timeouts to reduce flakiness in CI.

Author: C-Achard

dlclivegui/config.py

@@ -7,6 +7,8 @@
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
+from dlclivegui.temp import Engine
+
 Rotation = Literal[0, 90, 180, 270]
 TileLayout = Literal["auto", "2x2", "1x4", "4x1"]
 Precision = Literal["FP32", "FP16"]
@@ -239,7 +241,7 @@ class DLCProcessorSettings(BaseModel):
     resize: float = Field(default=1.0, gt=0)
     precision: Precision = "FP32"
     additional_options: dict[str, Any] = Field(default_factory=dict)
-    model_type: Literal["pytorch"] = "pytorch"
+    model_type: Engine = "pytorch"
     single_animal: bool = True
 
     @field_validator("dynamic", mode="before")

dlclivegui/gui/main_window.py

@@ -876,6 +876,14 @@ def _parse_json(self, value: str) -> dict:
 
     def _dlc_settings_from_ui(self) -> DLCProcessorSettings:
         model_path = self.model_path_edit.text().strip()
+        try:
+            DLCLiveProcessor.get_model_backend(model_path)
+        except Exception as e:
+            raise RuntimeError(
+                "Could not determine model backend from path."
+                "Please ensure the model file is valid and has an appropriate extension "
+                "(.pt, .pth for PyTorch or model directory for TensorFlow)."
+            ) from e
         return DLCProcessorSettings(
             model_path=model_path,
             model_directory=self._config.dlc.model_directory,  # Preserve from config
@@ -969,9 +977,9 @@ def _action_browse_model(self) -> None:
         dlg.setFileMode(QFileDialog.FileMode.ExistingFile)
         dlg.setNameFilters(
             [
-                "Model files (*.pt *.pth *.pb)",
+                "Model files (*.pt *.pth)",
                 "PyTorch models (*.pt *.pth)",
-                "TensorFlow models (*.pb)",
+                # "TensorFlow models (*.pb)",
                 "All files (*.*)",
             ]
         )

dlclivegui/services/dlc_processor.py

@@ -19,6 +19,7 @@
 
 # from dlclivegui.config import DLCProcessorSettings
 from dlclivegui.processors.processor_utils import instantiate_from_scan
+from dlclivegui.temp import Engine  # type: ignore # TODO use main package enum when released
 
 logger = logging.getLogger(__name__)
 
@@ -29,7 +30,6 @@
     from dlclive import (
         DLCLive,  # type: ignore
     )
-    from dlclivegui.temp import Engine  # type: ignore # TODO use main package one when released
 except Exception as e:  # pragma: no cover - handled gracefully
     logger.error(f"dlclive package could not be imported: {e}")
     DLCLive = None  # type: ignore[assignment]

dlclivegui/temp/__init__.py

@@ -0,0 +1,3 @@
+from engine import Engine  # type: ignore
+
+__all__ = ["Engine"]

dlclivegui/temp/engine.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from pathlib import Path
 
+
 class Engine(Enum):
     TENSORFLOW = "tensorflow"
     PYTORCH = "pytorch"
@@ -27,7 +28,7 @@ def from_model_path(cls, model_path: str | Path) -> "Engine":
             if has_cfg and has_pb:
                 return cls.TENSORFLOW
         elif path.is_file():
-            if path.suffix == ".pt":
+            if path.suffix in (".pt", ".pth"):
                 return cls.PYTORCH
 
-        raise ValueError(f"Could not determine engine from model path: {model_path}")
+        raise ValueError(f"Could not determine engine from model path: {model_path}")

tests/services/test_dlc_processor.py

@@ -63,10 +63,10 @@ def test_worker_processes_frames(qtbot, monkeypatch_dlclive, settings_model):
             proc.enqueue_frame(frame, timestamp=2.0 + i)
             qtbot.wait(5)  # ms
 
-        # NOTE @C-Achard this still fails randomly
-        # the timeout has to be surprisingly large here
-        # not sure if it's qtbot or threading scheduling delays
-        # Should be fixed now.
+        # NOTE @C-Achard The timeout here is intentionally large to account for potential
+        # Qt event-loop and threading scheduling delays in CI environments.
+        # This was previously flaky with a smaller timeout; increasing it should
+        # keep the test stable.
         qtbot.waitUntil(lambda: proc.get_stats().frames_processed >= 3, timeout=3000)
 
     finally: