mod_models.py

"""MOD model registry for the eve_gmod demo.

This is the single place that maps each selectable model name to the
``ModelConfig`` the EVE worker loads. Adding a model, or pointing an existing
one at new weights, is a one-line change here — nothing else needs to move.

- **GMOD Base Model** uses the EVE SDK's bundled default model, so it carries
  no path and no class table.
- **Automotive Object Detector** and **Office Object Detector** load their
  ``.tflite`` from ``MOD_MODELS_DIR``
  (the gated weights fetched at build time, see ``install_eve.py``; falls back to
  the demo-relative ``models/`` locally) and ship their own class labels (in
  output-index order).
"""

from __future__ import annotations

import os
from pathlib import Path

from eve_messages import ModelConfig

# In the Space the gated weights are fetched to MOD_MODELS_DIR by install_eve.py
# at Docker build time; local runs fall back to the demo-relative ``models/``.
_MODELS_DIR = Path(os.environ.get("MOD_MODELS_DIR") or Path(__file__).resolve().parent / "models")

# Class labels in the order the model emits them (output index -> label).
AUTOMOTIVE_CLASS_NAMES: tuple[str, ...] = (
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "bus",
    "truck",
    "traffic light",
    "stop sign",
)
OFFICE_CLASS_NAMES: tuple[str, ...] = (
    "bottle",
    "cup",
    "potted plant",
    "laptop",
    "mouse",
    "keyboard",
    "cell phone",
    "book",
)

MOD_MODEL_REGISTRY: dict[str, ModelConfig] = {
    "GMOD Base Model": ModelConfig(name="GMOD Base Model", nms_threshold=0.2),
    "Automotive Object Detector": ModelConfig(
        name="Automotive Object Detector",
        model_path=str(_MODELS_DIR / "automotive-640x640.tflite"),
        class_names=AUTOMOTIVE_CLASS_NAMES,
        nms_threshold=0.4,
        iou_threshold=0.4
    ),
    "Office Object Detector": ModelConfig(
        name="Office Object Detector",
        model_path=str(_MODELS_DIR / "office-640x640.tflite"),
        class_names=OFFICE_CLASS_NAMES,
        nms_threshold=0.4,
        iou_threshold=0.4
    ),
}

# UI order of the selection radio; first entry is the default.
MOD_MODELS: list[str] = list(MOD_MODEL_REGISTRY)
DEFAULT_MOD_MODEL: str = "GMOD Base Model"