Normalize model_type and fix model paths

C-Achard · C-Achard · commit 31e09d8852ee · 2026-02-23T11:52:30.000+01:00
Coerce DLCProcessorSettings.model_type to a lowercase string (accepting Enum or string inputs) and validate allowed backends (pytorch, tensorflow). Update UI to handle TensorFlow .pb models by using the parent directory for DLCLive, restrict file dialog to existing files, add existence checks and backend detection when selecting a model. Improve ModelPathStore: robust path normalization, separate helpers for existing file/dir checks, smarter save/load/resolve logic, and better start-dir/suggest-file heuristics. Minor cleanup: remove duplicate import and clarify a TF model detection comment in engine.
diff --git a/dlclivegui/config.py b/dlclivegui/config.py
@@ -2,13 +2,12 @@
 from __future__ import annotations
 
 import json
+from enum import Enum
 from pathlib import Path
 from typing import Any, Literal
 
 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"]
@@ -241,14 +240,46 @@ 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: Engine = Engine.PYTORCH
+    model_type: str = "pytorch"
     single_animal: bool = True
 
     @field_validator("dynamic", mode="before")
     @classmethod
     def _coerce_dynamic(cls, v):
         return DynamicCropModel.from_tupleish(v)
 
+    @field_validator("model_type", mode="before")
+    @classmethod
+    def _coerce_model_type(cls, v):
+        """
+        Accept:
+          - "pytorch"/"tensorflow"/etc as strings
+          - Enum instances (e.g. Engine.PYTORCH) and store their .value
+        Always return a lowercase string.
+        """
+        if v is None or v == "":
+            return "pytorch"
+
+        # If caller passed Engine enum or any Enum, use its value
+        if isinstance(v, Enum):
+            v = v.value
+
+        # If caller passed something with a `.value` attribute (defensive)
+        if not isinstance(v, str) and hasattr(v, "value"):
+            v = v.value
+
+        if not isinstance(v, str):
+            raise TypeError(f"model_type must be a string or Enum, got {type(v)!r}")
+
+        v = v.strip().lower()
+
+        # Optional: enforce allowed values
+        allowed = {"pytorch", "tensorflow"}
+        if v not in allowed:
+            raise ValueError(f"Unknown model type: {v!r}. Allowed: {sorted(allowed)}")
+
+        return v
+
 
 class BoundingBoxSettings(BaseModel):
     enabled: bool = False
diff --git a/dlclivegui/gui/main_window.py b/dlclivegui/gui/main_window.py
@@ -876,6 +876,10 @@ def _parse_json(self, value: str) -> dict:
 
     def _dlc_settings_from_ui(self) -> DLCProcessorSettings:
         model_path = self.model_path_edit.text().strip()
+        if Path(model_path).exists() and Path(model_path).suffix in (".pb"):
+            # IMPORTANT NOTE: DLClive expects a directory for TensorFlow models,
+            # so if user selects a .pb file, we should pass the parent directory to DLCLive
+            model_path = str(Path(model_path).parent)
         if model_path == "":
             raise ValueError("Model path cannot be empty. Please enter a valid path to a DLCLive model file.")
         try:
@@ -976,13 +980,12 @@ def _action_browse_model(self) -> None:
         preselect = self._model_path_store.suggest_selected_file()
 
         dlg = QFileDialog(self, "Select DLCLive model file")
-        dlg.setFileMode(QFileDialog.FileMode.AnyFile)
+        dlg.setFileMode(QFileDialog.FileMode.ExistingFile)
         dlg.setNameFilters(
             [
                 "Model files (*.pt *.pth)",
                 "PyTorch models (*.pt *.pth)",
-                # "TensorFlow models (*.pb)",
-                "TensorFlow model directory (*.*)",
+                "TensorFlow models (*.pb)",
             ]
         )
         dlg.setDirectory(start_dir)
@@ -995,7 +998,20 @@ def _action_browse_model(self) -> None:
             selected = dlg.selectedFiles()
             if not selected:
                 return
-            file_path = selected[0]
+            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
+
+            try:
+                DLCLiveProcessor.get_model_backend(str(file_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)
 
             # Persist model path + directory
diff --git a/dlclivegui/services/dlc_processor.py b/dlclivegui/services/dlc_processor.py
@@ -16,8 +16,6 @@
 from PySide6.QtCore import QObject, Signal
 
 from dlclivegui.config import DLCProcessorSettings
-
-# 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
 
diff --git a/dlclivegui/temp/engine.py b/dlclivegui/temp/engine.py
@@ -24,6 +24,7 @@ def from_model_path(cls, model_path: str | Path) -> "Engine":
 
         if path.is_dir():
             has_cfg = (path / "pose_cfg.yaml").is_file()
+            # has_cfg is DLClive specific and is considered a requirement for TF live models.
             has_pb = any(p.suffix == ".pb" for p in path.glob("*.pb"))
             if has_cfg and has_pb:
                 return cls.TENSORFLOW
diff --git a/dlclivegui/utils/settings_store.py b/dlclivegui/utils/settings_store.py
@@ -1,4 +1,6 @@
 # dlclivegui/utils/settings_store.py
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 
@@ -70,124 +72,194 @@ class ModelPathStore:
     def __init__(self, settings: QSettings | None = None):
         self._settings = settings or QSettings("DeepLabCut", "DLCLiveGUI")
 
-    def _norm(self, p: str | None) -> str | None:
+    # -------------------------
+    # Normalization helpers
+    # -------------------------
+    def _as_path(self, p: str | None) -> Path | None:
+        """Best-effort conversion to Path (expand ~, interpret '.' as cwd)."""
         if not p:
             return None
+        s = str(p).strip()
+        if not s:
+            return None
         try:
-            return str(Path(p).expanduser().resolve())
+            pp = Path(s).expanduser()
+            if s in (".", "./"):
+                pp = Path.cwd()
+            return pp
         except Exception:
-            logger.debug("Failed to normalize path: %s", p)
+            logger.debug("Failed to parse path: %s", p)
+            return None
+
+    def _norm_existing_dir(self, p: str | None) -> str | None:
+        """Return an absolute, resolved existing directory path, else None."""
+        pp = self._as_path(p)
+        if pp is None:
+            return None
+        try:
+            # If a file was given, use its parent directory
+            if pp.exists() and pp.is_file():
+                pp = pp.parent
+
+            if pp.exists() and pp.is_dir():
+                return str(pp.resolve())
+        except Exception:
+            logger.debug("Failed to normalize directory: %s", p)
+        return None
+
+    def _norm_existing_path(self, p: str | None) -> str | None:
+        """Return an absolute, resolved existing path (file or dir), else None."""
+        pp = self._as_path(p)
+        if pp is None:
             return None
+        try:
+            if pp.exists():
+                return str(pp.resolve())
+        except Exception:
+            logger.debug("Failed to normalize path: %s", p)
+        return None
 
+    # -------------------------
+    # Load
+    # -------------------------
     def load_last(self) -> str | None:
+        """Return last model path if it still exists and looks usable."""
         val = self._settings.value("dlc/last_model_path")
-        path = self._norm(str(val)) if val else None
+        path = self._norm_existing_path(str(val)) if val else None
         if not path:
             return None
+
         try:
-            return path if is_model_file(path) else None
+            pp = Path(path)
+            # Accept a valid model *file*
+            if pp.is_file() and is_model_file(str(pp)):
+                return str(pp)
         except Exception:
-            logger.debug("Last model path is not a valid model file: %s", path)
-            return None
+            logger.debug("Last model path not valid/usable: %s", path)
+
+        return None
 
     def load_last_dir(self) -> str | None:
+        """Return last directory if it still exists and is a directory."""
         val = self._settings.value("dlc/last_model_dir")
-        d = self._norm(str(val)) if val else None
-        if not d:
-            return None
-        try:
-            p = Path(d)
-            return str(p) if p.exists() and p.is_dir() else None
-        except Exception:
-            logger.debug("Last model dir is not a valid directory: %s", d)
-            return None
+        d = self._norm_existing_dir(str(val)) if val else None
+        return d
 
+    # -------------------------
+    # Save
+    # -------------------------
     def save_if_valid(self, path: str) -> None:
-        """Save last model *file* if it looks valid, and always save its directory."""
-        path = self._norm(path) or ""
-        if not path:
+        """
+        Save last model path if it looks valid/usable, and always save its directory.
+        - For files: always save parent directory.
+        - For directories: save directory itself if it looks like a TF model dir.
+        """
+        norm = self._norm_existing_path(path)
+        if not norm:
             return
+
         try:
-            parent = str(Path(path).parent)
-            self._settings.setValue("dlc/last_model_dir", parent)
+            p = Path(norm)
+
+            # Always persist a *directory* that is safe for QFileDialog.setDirectory(...)
+            if p.is_dir():
+                model_dir = p
+            else:
+                model_dir = p.parent
+
+            model_dir_norm = self._norm_existing_dir(str(model_dir))
+            if model_dir_norm:
+                self._settings.setValue("dlc/last_model_dir", model_dir_norm)
+
+            # Persist model path if it is a valid model file, or a TF model directory
+            if p.is_file() and is_model_file(str(p)):
+                self._settings.setValue("dlc/last_model_path", str(p))
+            elif p.is_dir() and self._looks_like_tf_model_dir(p):
+                self._settings.setValue("dlc/last_model_path", str(p))
 
-            if is_model_file(path):
-                self._settings.setValue("dlc/last_model_path", str(Path(path)))
         except Exception:
-            logger.debug("Failed to save last model path: %s", path)
-            pass
+            logger.debug("Failed to save model path: %s", path, exc_info=True)
 
     def save_last_dir(self, directory: str) -> None:
-        directory = self._norm(directory) or ""
-        if not directory:
+        d = self._norm_existing_dir(directory)
+        if not d:
             return
         try:
-            p = Path(directory)
-            if p.exists() and p.is_dir():
-                self._settings.setValue("dlc/last_model_dir", str(p))
+            self._settings.setValue("dlc/last_model_dir", d)
         except Exception:
-            pass
+            logger.debug("Failed to save last model dir: %s", d, exc_info=True)
 
+    # -------------------------
+    # Resolve
+    # -------------------------
     def resolve(self, config_path: str | None) -> str:
-        """Resolve the best model path to display in the UI."""
-        config_path = self._norm(config_path)
-        if config_path:
+        """
+        Resolve the best model path to display in the UI.
+        Preference:
+          1) config_path if valid/usable
+          2) persisted last model path if valid/usable
+          3) empty
+        """
+        cfg = self._norm_existing_path(config_path)
+        if cfg:
             try:
-                if is_model_file(config_path):
-                    return config_path
+                p = Path(cfg)
+                if p.is_file() and is_model_file(cfg):
+                    return cfg
+                if p.is_dir() and self._looks_like_tf_model_dir(p):
+                    return cfg
             except Exception:
-                logger.debug("Config path is not a valid model file: %s", config_path)
-                pass
+                logger.debug("Config path not usable: %s", cfg)
 
         persisted = self.load_last()
         if persisted:
-            try:
-                if is_model_file(persisted):
-                    return persisted
-            except Exception:
-                pass
+            return persisted
 
         return ""
 
     def suggest_start_dir(self, fallback_dir: str | None = None) -> str:
-        """Pick the best directory to start the file dialog in."""
+        """
+        Pick the best directory to start file dialogs in.
+        Guarantees: returns an existing absolute directory (never '.').
+        """
         # 1) last dir
         last_dir = self.load_last_dir()
         if last_dir:
             return last_dir
 
-        # 2) directory of last valid model file
-        last_file = self.load_last()
-        if last_file:
+        # 2) directory of last valid model path
+        last = self.load_last()
+        if last:
             try:
-                parent = Path(last_file).parent
-                if parent.exists():
-                    return str(parent)
+                p = Path(last)
+                if p.is_file():
+                    parent = self._norm_existing_dir(str(p.parent))
+                    if parent:
+                        return parent
+                elif p.is_dir():
+                    d = self._norm_existing_dir(str(p))
+                    if d:
+                        return d
             except Exception:
-                logger.debug("Failed to get parent of last model file: %s", last_file)
-                pass
+                logger.debug("Failed to derive start dir from last model: %s", last)
 
-        # 3) fallback dir (config.model_directory) if valid
-        if fallback_dir:
-            try:
-                p = Path(fallback_dir).expanduser()
-                if p.exists() and p.is_dir():
-                    return str(p)
-            except Exception:
-                logger.debug("Fallback dir is not a valid directory: %s", fallback_dir)
-                pass
+        # 3) fallback dir (e.g. config.dlc.model_directory)
+        fb = self._norm_existing_dir(fallback_dir)
+        if fb:
+            return fb
 
-        # 4) last resort: home
-        return str(Path.home())
+        # 4) last resort: cwd if exists else home
+        cwd = self._norm_existing_dir(str(Path.cwd()))
+        return cwd or str(Path.home())
 
     def suggest_selected_file(self) -> str | None:
-        """Optional: return a file to preselect if it exists."""
-        last_file = self.load_last()
-        if not last_file:
+        """Return a file to preselect if it exists (only files, not directories)."""
+        last = self.load_last()
+        if not last:
             return None
         try:
-            p = Path(last_file)
+            p = Path(last)
             return str(p) if p.exists() and p.is_file() else None
         except Exception:
-            logger.debug("Failed to check existence of last model file: %s", last_file)
+            logger.debug("Failed to check existence of last model: %s", last)
             return None
