endoreg-db 0.8.6.1__py3-none-any.whl → 0.8.8.9__py3-none-any.whl

endoreg_db/management/commands/create_multilabel_model_meta.py

@@ -1,214 +1,384 @@
 """
-This script creates a new ModelMeta object for a multilabel classification model.
+Management command for creating ModelMeta entries for multilabel classification models.
+
+Supports two workflows:
+1. Registering a local `.safetensors` weights file.
+2. Generating metadata from a YAML template, downloading weights from Hugging Face.
 """
 
+import logging
+import tempfile
 from pathlib import Path
-from django.core.management import BaseCommand
-from endoreg_db.models import LabelSet, AiModel, ModelMeta
+from typing import Any, Dict, Iterable, List
+
+import yaml
+from django.core.management import BaseCommand, CommandError
+from huggingface_hub import hf_hub_download
 
+from endoreg_db.data import AI_MODEL_META_DATA_DIR
+from endoreg_db.models import AiModel, LabelSet, ModelMeta
 
-# Example usage:
-# python manage.py create_multilabel_model_meta --model_path "./data/models/colo_segmentation_RegNetX800MF_6.ckpt"
-# python manage.py create_multilabel_model_meta --model_path "./libs/endoreg-db/tests/assets/colo_segmentation_RegNetX800MF_6.ckpt"
+logger = logging.getLogger(__name__)
 
-FPS = 50
-SMOOTH_WINDOW_SIZE_S = 1
-MIN_SEQ_LEN_S = 0.5
-crop_template = [0, 1080, 550, 1920 - 20]  # [top, bottom, left, right]
+DEFAULT_MODEL_PATH = (
+    Path(__file__).resolve().parents[3]
+    / "tests"
+    / "assets"
+    / "colo_segmentation_RegNetX800MF_6.safetensors"
+)
 
 
 class Command(BaseCommand):
-    help = """
-        Creates a new ModelMeta object for a multilabel classification model.
-        1. Validates the existence of the specified LabelSet and AiModel
-        2. Checks that the model file exists on disk
-        3. Creates or updates a ModelMeta entry with the specified parameters
-    """
-
-    def add_arguments(self, parser):
-        """
-        Adds command-line arguments for configuring model meta creation.
-        
-        This method defines all necessary arguments for specifying model details, label set selection, normalization parameters, image dimensions, axes configuration, batch size, worker count, and versioning options for the Django management command.
-        """
+    help = (
+        "Create or update ModelMeta entries for multilabel classification models using "
+        "either a local safetensor file or a YAML template with Hugging Face download support."
+    )
+
+    def add_arguments(self, parser):  # noqa: D401 - inherited docstring is sufficient
         parser.add_argument(
             "--model_name",
             type=str,
             default="image_multilabel_classification_colonoscopy_default",
-            help="Name of the model to use for segmentation",
+            help="Name of the AiModel to attach metadata to.",
         )
-
-        # model_path
         parser.add_argument(
             "--model_path",
             type=str,
-            #default="./data/models/colo_segmentation_RegNetX800MF_6.ckpt", #model not yet here
-            default="./tests/assets/colo_segmentation_RegNetX800MF_6.ckpt", # model currently still here
-            help="Path to the model file",
+            default=str(DEFAULT_MODEL_PATH),
+            help=(
+                "Path to a local .safetensors weights file. If provided (or left as default) "
+                "the command registers the local weights."
+            ),
+        )
+        parser.add_argument(
+            "--template_path",
+            type=str,
+            default=None,
+            help="Absolute or relative path to a model meta YAML template.",
+        )
+        parser.add_argument(
+            "--template_name",
+            type=str,
+            default=None,
+            help=(
+                "Name of a built-in template file in endoreg_db/data/ai_model_meta (without extension)."
+            ),
+        )
+        parser.add_argument(
+            "--template_entry_name",
+            type=str,
+            default=None,
+            help="Optional entry selector when the template file defines multiple models.",
         )
-
-        # model_meta_version: int = 1
         parser.add_argument(
             "--model_meta_version",
-            type=int,
-            default=1,
-            help="Version of the model meta",
+            type=str,
+            default=None,
+            help=(
+                "Version to assign to the metadata. When omitted the command uses the template value "
+                "or defaults to '1' for local registrations."
+            ),
         )
-
         parser.add_argument(
             "--image_classification_labelset_name",
             type=str,
             default="multilabel_classification_colonoscopy_default",
-            help="Name of the LabelSet for image classification",
+            help="Name of the LabelSet used by the model.",
         )
-
         parser.add_argument(
             "--image_classification_labelset_version",
             type=int,
             default=-1,
-            help="Version of the LabelSet for image classification, default is -1 which invokes the highest version",
+            help="Specific LabelSet version. Use -1 to select the latest available version.",
         )
-
-        # activation: str = "sigmoid"
         parser.add_argument(
             "--activation_function_name",
             type=str,
             default="sigmoid",
-            help="Activation function for the model",
+            help="Activation function applied to model outputs.",
         )
-
-        # mean: str = "0.45211223,0.27139644,0.19264949"
         parser.add_argument(
             "--mean",
             type=str,
             default="0.45211223,0.27139644,0.19264949",
-            help="Mean for normalization",
+            help="Comma-separated mean values for input normalization.",
         )
-
         parser.add_argument(
             "--std",
             type=str,
             default="0.31418097,0.21088019,0.16059452",
-            help="Std for normalization",
+            help="Comma-separated std values for input normalization.",
         )
-
-        # size_x: int = 716
         parser.add_argument(
             "--size_x",
             type=int,
             default=716,
-            help="Size of the image in x direction",
+            help="Input width expected by the model.",
         )
-
-        # size_y: int = 716
         parser.add_argument(
             "--size_y",
             type=int,
             default=716,
-            help="Size of the image in y direction",
+            help="Input height expected by the model.",
         )
-
         parser.add_argument(
             "--axes",
             type=str,
             default="2,0,1",
-            help="Axes of the image",
+            help="Comma-separated axis order expected by the model (e.g. '2,0,1' for CHW).",
         )
-
-        # batchsize: int = 16
         parser.add_argument(
             "--batchsize",
             type=int,
             default=16,
-            help="Batchsize for the model",
+            help="Default batch size for inference.",
         )
-
-        # num_workers: int = 0
         parser.add_argument(
             "--num_workers",
             type=int,
             default=0,
-            help="Number of workers for the model",
+            help="Default number of data loading workers.",
+        )
+        parser.add_argument(
+            "--description",
+            type=str,
+            default="",
+            help="Optional description to store on the ModelMeta record.",
         )
-
-        # bump_version: bool = False
         parser.add_argument(
             "--bump_version",
             action="store_true",
-            help="Bump the version of the model",
+            help="If the requested version exists, bump to the next available version instead of failing.",
+        )
+        parser.add_argument(
+            "--huggingface_token",
+            type=str,
+            default=None,
+            help="Optional Hugging Face token for private repositories.",
+        )
+
+    def handle(self, *args, **options):  # noqa: D401 - inherited docstring is sufficient
+        use_template = options.get("template_path") or options.get("template_name")
+
+        try:
+            if use_template:
+                model_meta = self._create_from_template(options)
+            else:
+                model_meta = self._create_from_local_file(options)
+        except CommandError:
+            raise
+        except Exception as exc:  # pragma: no cover - defensive logging
+            logger.exception("Failed to create ModelMeta", exc_info=exc)
+            raise CommandError(str(exc)) from exc
+
+        self.stdout.write(
+            self.style.SUCCESS(f"ModelMeta ready: {model_meta.name} (v{model_meta.version}) for {model_meta.model.name}")
         )
 
-    def handle(self, *args, **options):
-        """
-        Processes command-line arguments to create or update a ModelMeta object for a multilabel classification model.
-        
-        Retrieves and validates the specified LabelSet and AiModel, ensures the model file exists, and invokes ModelMeta.create_from_file with the provided configuration. Raises a ValueError if the requested LabelSet does not exist.
-        """
+    def _create_from_local_file(self, options: Dict[str, Any]) -> ModelMeta:
+        weights_path = Path(options["model_path"]).expanduser().resolve()
+        self._validate_safetensors_path(weights_path)
+
         model_name = options["model_name"]
+        self._ensure_ai_model_exists(model_name)
+
+        labelset = self._resolve_labelset(
+            options["image_classification_labelset_name"],
+            options.get("image_classification_labelset_version"),
+        )
 
-        model_path = options["model_path"]
-        activation_function_name = options["activation_function_name"]
+        requested_version = options.get("model_meta_version") or "1"
 
-        mean = options["mean"]
-        std = options["std"]
+        model_meta = ModelMeta.create_from_file(
+            meta_name=model_name,
+            model_name=model_name,
+            labelset_name=labelset.name,
+            labelset_version=labelset.version,
+            weights_file=weights_path.as_posix(),
+            requested_version=str(requested_version),
+            bump_if_exists=options.get("bump_version", False),
+            **self._collect_local_kwargs(options),
+        )
 
-        size_x = options["size_x"]
-        size_y = options["size_y"]
+        return model_meta
 
-        axes = options["axes"]
+    def _create_from_template(self, options: Dict[str, Any]) -> ModelMeta:
+        template_path = self._resolve_template_path(options)
+        entries = self._load_template_entries(template_path)
+        entry = self._select_template_entry(entries, options)
 
-        image_classification_labelset_name = options[
-            "image_classification_labelset_name"
-        ]
+        fields = entry.get("fields", {})
+        if not fields:
+            raise CommandError("Template entry is missing a 'fields' section.")
 
-        image_classification_labelset_version = options[
-            "image_classification_labelset_version"
-        ]
-        if image_classification_labelset_version == -1:
-            # If version is -1, we want the latest version of the labelset
-            labelset = LabelSet.objects.filter(
-                name=image_classification_labelset_name
-            ).order_by("-version").first()
-            if labelset:
-                image_classification_labelset_version = labelset.version
-            else:
-                raise ValueError(
-                    f"No LabelSet found with name: {image_classification_labelset_name}"
-                )
-        else:
-            # If a specific version is provided, we use that
-            labelset = LabelSet.objects.filter(
-                name=image_classification_labelset_name,
-                version=image_classification_labelset_version
-            ).first()
-            if not labelset:
-                raise ValueError(
-                    f"No LabelSet found with name: {image_classification_labelset_name} and version: {image_classification_labelset_version}"
-                )
-
-        # assert model exists
-        db_ai_model = AiModel.objects.filter(name=model_name).first()
-        assert db_ai_model, f"Model not found: {model_name}"
-
-        model_path = Path(model_path).expanduser().resolve().as_posix()
-
-        assert Path(model_path).exists(), f"Model file not found: {model_path}"
-
-        _model_meta = ModelMeta.create_from_file(
-            meta_name=model_name,  # Name of the new metadata object
-            model_name=model_name,  # name of the ai_model to use; This also defines the models VideoSegmentationLabelSet
-            labelset_name=image_classification_labelset_name,
-            weights_file=model_path,
-            # Use the correct keyword arguments matching the method signature
-            requested_version=options["model_meta_version"],
-            bump_if_exists=options["bump_version"],
-            # Pass other options via **kwargs
-            activation=activation_function_name,
-            mean=mean,
-            std=std,
-            size_x=size_x,
-            size_y=size_y,
-            axes=axes,
-            batchsize=options["batchsize"],
-            num_workers=options["num_workers"],
+        meta_name = fields.get("name") or options["model_name"]
+        model_name = fields.get("model") or options["model_name"]
+        labelset_name = fields.get("labelset") or options["image_classification_labelset_name"]
+        labelset_version = fields.get("labelset_version", options.get("image_classification_labelset_version"))
+
+        self._ensure_ai_model_exists(model_name)
+        labelset = self._resolve_labelset(labelset_name, labelset_version)
+
+        requested_version = options.get("model_meta_version") or fields.get("version")
+        if not requested_version:
+            raise CommandError("Provide --model_meta_version or include a 'version' in the template entry.")
+
+        hf_config = entry.get("setup_config", {}).get("huggingface_fallback", {})
+        repo_id = hf_config.get("repo_id")
+        filename = hf_config.get("filename")
+
+        if not repo_id or not filename:
+            raise CommandError(
+                "Template entry must define setup_config.huggingface_fallback.repo_id and filename for weight download."
+            )
+
+        if not filename.endswith(".safetensors"):
+            raise CommandError("Only .safetensors files are supported when downloading from Hugging Face.")
+
+        token = options.get("huggingface_token")
+
+        with tempfile.TemporaryDirectory(prefix="hf-multilabel-") as download_dir:
+            download_kwargs = {
+                "repo_id": repo_id,
+                "filename": filename,
+                "local_dir": download_dir,
+                "local_dir_use_symlinks": False,
+            }
+            if token:
+                download_kwargs["token"] = token
+
+            weights_path = Path(hf_hub_download(**download_kwargs)).resolve()
+
+            self._validate_safetensors_path(weights_path)
+
+            model_meta = ModelMeta.create_from_file(
+                meta_name=meta_name,
+                model_name=model_name,
+                labelset_name=labelset.name,
+                labelset_version=labelset.version,
+                weights_file=weights_path.as_posix(),
+                requested_version=str(requested_version),
+                bump_if_exists=options.get("bump_version", False),
+                **self._collect_template_kwargs(fields, options),
+            )
+
+        return model_meta
+
+    def _resolve_template_path(self, options: Dict[str, Any]) -> Path:
+        template_path = options.get("template_path")
+        template_name = options.get("template_name")
+
+        if template_path:
+            resolved = Path(template_path).expanduser().resolve()
+        elif template_name:
+            resolved = (AI_MODEL_META_DATA_DIR / f"{template_name}.yaml").resolve()
+        else:  # pragma: no cover - guarded by caller
+            raise CommandError("Template mode requires --template_path or --template_name.")
+
+        if not resolved.exists():
+            raise CommandError(f"Template file not found: {resolved}")
+
+        return resolved
+
+    @staticmethod
+    def _load_template_entries(template_path: Path) -> List[Dict[str, Any]]:
+        with template_path.open("r", encoding="utf-8") as handle:
+            data = yaml.safe_load(handle) or []
+
+        if isinstance(data, dict):
+            return [data]
+        if isinstance(data, list):
+            return [entry for entry in data if isinstance(entry, dict)]
+
+        raise CommandError(f"Template {template_path} must define a mapping or list of mappings.")
+
+    def _select_template_entry(self, entries: Iterable[Dict[str, Any]], options: Dict[str, Any]) -> Dict[str, Any]:
+        target = options.get("template_entry_name") or options.get("model_name")
+
+        for entry in entries:
+            fields = entry.get("fields", {})
+            if not fields:
+                continue
+            if target and (fields.get("name") == target or fields.get("model") == target):
+                return entry
+
+        entries = list(entries)
+        if len(entries) == 1:
+            return entries[0]
+
+        raise CommandError(
+            "Unable to determine which template entry to use. Specify --template_entry_name to disambiguate."
+        )
+
+    def _collect_local_kwargs(self, options: Dict[str, Any]) -> Dict[str, Any]:
+        return self._filter_none(
+            {
+                "activation": options.get("activation_function_name"),
+                "mean": options.get("mean"),
+                "std": options.get("std"),
+                "size_x": options.get("size_x"),
+                "size_y": options.get("size_y"),
+                "axes": options.get("axes"),
+                "batchsize": options.get("batchsize"),
+                "num_workers": options.get("num_workers"),
+                "description": options.get("description"),
+            }
+        )
+
+    def _collect_template_kwargs(self, fields: Dict[str, Any], options: Dict[str, Any]) -> Dict[str, Any]:
+        def numeric(value):
+            return int(value) if value is not None else value
+
+        return self._filter_none(
+            {
+                "activation": fields.get("activation") or options.get("activation_function_name"),
+                "mean": self._normalise_sequence(fields.get("mean")) or options.get("mean"),
+                "std": self._normalise_sequence(fields.get("std")) or options.get("std"),
+                "size_x": numeric(fields.get("size_x")) if fields.get("size_x") is not None else options.get("size_x"),
+                "size_y": numeric(fields.get("size_y")) if fields.get("size_y") is not None else options.get("size_y"),
+                "axes": fields.get("axes") or options.get("axes"),
+                "batchsize": numeric(fields.get("batchsize")) if fields.get("batchsize") is not None else options.get("batchsize"),
+                "num_workers": numeric(fields.get("num_workers")) if fields.get("num_workers") is not None else options.get("num_workers"),
+                "description": fields.get("description") or options.get("description"),
+            }
         )
+
+    @staticmethod
+    def _normalise_sequence(value: Any) -> str | None:
+        if value is None:
+            return None
+        if isinstance(value, str):
+            return value
+        if isinstance(value, (list, tuple)):
+            return ",".join(str(item) for item in value)
+        return str(value)
+
+    @staticmethod
+    def _filter_none(payload: Dict[str, Any]) -> Dict[str, Any]:
+        return {key: value for key, value in payload.items() if value not in (None, "")}
+
+    @staticmethod
+    def _validate_safetensors_path(path: Path) -> None:
+        if path.suffix != ".safetensors":
+            raise CommandError(f"Expected a .safetensors file, got: {path}")
+        if not path.exists():
+            raise CommandError(f"Weights file not found: {path}")
+
+    @staticmethod
+    def _ensure_ai_model_exists(model_name: str) -> None:
+        if not AiModel.objects.filter(name=model_name).exists():
+            raise CommandError(f"AiModel not found: {model_name}. Load ai model data before running this command.")
+
+    @staticmethod
+    def _resolve_labelset(name: str, version: Any) -> LabelSet:
+        queryset = LabelSet.objects.filter(name=name)
+
+        if version in (None, -1):
+            labelset = queryset.order_by("-version").first()
+        else:
+            labelset = queryset.filter(version=version).first()
+
+        if not labelset:
+            raise CommandError(f"LabelSet not found for name='{name}' and version='{version}'.")
+
+        return labelset