Automatically log annotation Specs (DAT-1067)

datamint/api/endpoints/models_api.py

@@ -13,7 +13,13 @@
 from .model_types import Model, ModelVersion
 
 if TYPE_CHECKING:
+    from mlflow.models import ModelInputExample, ModelSignature
+
+    from datamint.dataset.base import DatamintBaseDataset
+    from datamint.entities.annotations.annotation_spec import AnnotationSpec
     from datamint.entities.project import Project
+    from datamint.mlflow.flavors.model import BaseDatamintModel
+    from datamint.mlflow.flavors.task_type import TaskType
 
     from .projects_api import ProjectsApi
 
@@ -244,3 +250,93 @@ def clone_model(self,
                 set_project(previous_project_id)
             else:
                 _reset_active_project()
+
+    def log_model(self,
+                  datamint_model: 'BaseDatamintModel',
+                  project: 'str | Project',
+                  model_name: str,
+                  *,
+                  dataset: 'DatamintBaseDataset | None' = None,
+                  task_type: 'TaskType | str | None' = None,
+                  annotation_specs: 'Sequence[AnnotationSpec] | None' = None,
+                  supported_modes: 'Sequence[str] | None' = None,
+                  code_paths: 'Sequence[str] | None' = None,
+                  artifacts: dict | None = None,
+                  signature: 'ModelSignature | None' = None,
+                  input_example: 'ModelInputExample | None' = None,
+                  pip_requirements: 'Sequence[str] | None' = None,
+                  extra_pip_requirements: 'Sequence[str] | None' = None,
+                  ) -> Model:
+        """Log a model to the Datamint model registry, opening its own MLflow run.
+
+        A thin wrapper around :func:`~datamint.mlflow.flavors.datamint_flavor.log_model` 
+        for models trained outside a Datamint trainer (the trainers log their models
+        automatically). It resolves the active project, starts an MLflow run,
+        and registers the model under ``model_name`` -- same run-handling as
+        :meth:`clone_model`.
+
+        Args:
+            datamint_model: The trained model to log.
+            project: Project (name, ID, or :class:`~datamint.entities.project.Project`)
+                to register the model under.
+            model_name: Name to register the model under. Also used as the
+                MLflow experiment name.
+            dataset: The dataset the model was trained on, used to derive
+                ``annotation_specs`` automatically when they aren't given
+                explicitly. Same dataset object you already have from training
+                (e.g. an :class:`~datamint.dataset.ImageDataset`).
+            task_type: Task type used to pick the right annotation-spec
+                builder. Defaults to ``datamint_model.task_type`` when omitted.
+            annotation_specs: Explicit annotation specs. Takes precedence over
+                anything derived from ``dataset``.
+            supported_modes: Prediction modes the model supports.
+            code_paths: Local paths to custom code files/dirs the model's
+                class depends on.
+            artifacts: Extra artifacts to bundle with the model.
+            signature: MLflow model signature.
+            input_example: MLflow input example.
+            pip_requirements: Exact pip requirements for the model environment.
+            extra_pip_requirements: Additional pip requirements on top of the
+                inferred defaults.
+
+        Returns:
+            The newly registered :class:`~.model_types.Model`.
+        """
+        from datamint.mlflow.flavors import datamint_flavor
+        from datamint.mlflow.flavors.annotation_specs import build_annotation_specs_for_task
+        from datamint.mlflow.tracking.fluent import (
+            _reset_active_project,
+            get_active_project_id,
+            set_project,
+        )
+
+        resolved_task_type = task_type or getattr(datamint_model, 'task_type', None)
+
+        resolved_specs = annotation_specs
+        if resolved_specs is None and dataset is not None:
+            resolved_specs = build_annotation_specs_for_task(resolved_task_type, dataset)
+
+        previous_project_id = get_active_project_id()
+        try:
+            set_project(project)
+            mlflow.set_experiment(model_name)
+            with mlflow.start_run(run_name=f"log_{model_name}"):
+                datamint_flavor.log_model(
+                    datamint_model,
+                    task_type=resolved_task_type,
+                    supported_modes=supported_modes,
+                    annotation_specs=resolved_specs,
+                    model_name=model_name,
+                    code_paths=code_paths,
+                    artifacts=artifacts,
+                    signature=signature,
+                    input_example=input_example,
+                    pip_requirements=pip_requirements,
+                    extra_pip_requirements=extra_pip_requirements,
+                )
+            return self.get_by_name(model_name)
+        finally:
+            if previous_project_id is not None:
+                set_project(previous_project_id)
+            else:
+                _reset_active_project()

datamint/lightning/trainers/classification_trainer.py

@@ -10,7 +10,7 @@
 
 from datamint.dataset import ImageDataset
 from datamint.entities.annotations.annotation_spec import CategoryAnnotationSpec
-from datamint.entities.annotations.types import AnnotationType
+from datamint.mlflow.flavors.annotation_specs import build_classification_annotation_specs
 
 from .base_trainer import BaseTrainer
 from .lightning_modules import ClassificationModule
@@ -47,19 +47,7 @@ def _monitor_metric(self) -> tuple[str, str]:
         return 'val/accuracy', 'max'
 
     def _build_annotation_specs(self) -> list[CategoryAnnotationSpec]:
-        groups: dict[str, list[str]] = {}
-        for identifier, value in self.dataset.image_categories_set:
-            groups.setdefault(identifier, []).append(value)
-        return [
-            CategoryAnnotationSpec(
-                type=AnnotationType.CATEGORY,
-                scope='image',
-                identifier=ident,
-                required=True,
-                values=sorted(vals),
-            )
-            for ident, vals in groups.items()
-        ]
+        return build_classification_annotation_specs(self.dataset)
 
 
 class ImageClassificationTrainer(ClassificationTrainer):

datamint/lightning/trainers/detection_trainer.py

@@ -6,8 +6,8 @@
 
 from datamint.dataset.image_dataset import ImageDataset, detection_collate_fn
 from datamint.entities.annotations.annotation_spec import AnnotationSpec
-from datamint.entities.annotations.types import AnnotationType
 from datamint.lightning.datamodule import DatamintDataModule
+from datamint.mlflow.flavors.annotation_specs import build_detection_annotation_specs
 
 from .base_trainer import BaseTrainer
 
@@ -67,7 +67,4 @@ def _monitor_metric(self) -> tuple[str, str]:
         return 'val/map', 'max'
 
     def _build_annotation_specs(self) -> list[AnnotationSpec]:
-        return [
-            AnnotationSpec(type=AnnotationType.SQUARE, scope='image', identifier=name, required=False)
-            for name in self.dataset.box_labels_set
-        ]
+        return build_detection_annotation_specs(self.dataset)

datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py

@@ -20,6 +20,8 @@
 from torch import Tensor, nn
 from typing_extensions import override
 
+from datamint.mlflow.flavors.task_type import TaskType
+
 from ..segmentation_module import SegmentationModule
 
 # ---------------------------------------------------------------------------
@@ -417,6 +419,7 @@ def forward(self, x: Tensor) -> Tensor:
 # ---------------------------------------------------------------------------
 
 class UNETRPPModule(SegmentationModule):
+    task_type = TaskType.VOLUME_SEGMENTATION
     """Segmentation module wrapping UNETR++ for 3-D volumetric inputs.
 
     Differences from 2-D modules:

datamint/lightning/trainers/segmentation_trainer.py

@@ -9,7 +9,7 @@
 from torch import nn
 
 from datamint.entities.annotations.annotation_spec import AnnotationSpec
-from datamint.entities.annotations.types import AnnotationType
+from datamint.mlflow.flavors.annotation_specs import build_segmentation_annotation_specs
 
 from .base_trainer import BaseTrainer
 
@@ -61,7 +61,4 @@ def _monitor_metric(self) -> tuple[str, str]:
         return 'val/iou', 'max'
 
     def _build_annotation_specs(self) -> list[AnnotationSpec]:
-        return [
-            AnnotationSpec(type=AnnotationType.SEGMENTATION, scope='image', identifier=name, required=False)
-            for name in self.dataset.seglabel_list
-        ]
+        return build_segmentation_annotation_specs(self.dataset, scope='image')

datamint/lightning/trainers/specialized/nnunet/trainer.py

@@ -12,8 +12,8 @@
 
 from datamint.dataset.volume_dataset import VolumeDataset
 from datamint.entities.annotations.annotation_spec import AnnotationSpec
-from datamint.entities.annotations.types import AnnotationType
 from datamint.lightning.trainers.base_trainer import BaseTrainer
+from datamint.mlflow.flavors.annotation_specs import build_segmentation_annotation_specs
 from datamint.lightning.trainers.specialized.nnunet.data_export import (
     DatamintToNNUNetExporter,
 )
@@ -113,10 +113,7 @@ def _build_annotation_specs(self) -> list[AnnotationSpec]:
             allow_external_annotations=True,
             include_unannotated=False,
         )
-        return [
-            AnnotationSpec(type=AnnotationType.SEGMENTATION, scope='volume', identifier=name, required=False)
-            for name in ds.seglabel_list
-        ]
+        return build_segmentation_annotation_specs(ds, scope='volume')
 
     def _build_model(self, *args, **kwargs):
         raise NotImplementedError(

datamint/lightning/trainers/vol_seg_trainer.py

@@ -13,8 +13,8 @@
 
 from datamint.dataset import VolumeDataset
 from datamint.entities.annotations.annotation_spec import AnnotationSpec
-from datamint.entities.annotations.types import AnnotationType
 from datamint.lightning.datamodule import DatamintDataModule
+from datamint.mlflow.flavors.annotation_specs import build_segmentation_annotation_specs
 
 from .segmentation_trainer import SegmentationTrainer
 
@@ -166,7 +166,4 @@ def _build_datamodule(
         )
 
     def _build_annotation_specs(self) -> list[AnnotationSpec]:
-        return [
-            AnnotationSpec(type=AnnotationType.SEGMENTATION, scope='volume', identifier=name, required=False)
-            for name in self.dataset.seglabel_list
-        ]
+        return build_segmentation_annotation_specs(self.dataset, scope='volume')

datamint/mlflow/flavors/annotation_specs.py

@@ -0,0 +1,72 @@
+"""Derive AnnotationSpec lists from a trained dataset's actual labels. """
+
+from typing import TYPE_CHECKING
+
+from datamint.entities.annotations.annotation_spec import AnnotationSpec, CategoryAnnotationSpec
+from datamint.entities.annotations.types import AnnotationType
+
+from .task_type import TaskType
+
+if TYPE_CHECKING:
+    from datamint.dataset.base import DatamintBaseDataset
+
+
+def build_segmentation_annotation_specs(dataset: 'DatamintBaseDataset',
+                                        scope: str = 'image') -> list[AnnotationSpec]:
+    """Build one AnnotationSpec per segmentation label in the dataset."""
+    return [
+        AnnotationSpec(type=AnnotationType.SEGMENTATION, scope=scope, identifier=name, required=False)
+        for name in dataset.seglabel_list
+    ]
+
+
+def build_detection_annotation_specs(dataset: 'DatamintBaseDataset') -> list[AnnotationSpec]:
+    """Build one AnnotationSpec per box-annotation class in the dataset."""
+    return [
+        AnnotationSpec(type=AnnotationType.SQUARE, scope='image', identifier=name, required=False)
+        for name in dataset.box_labels_set
+    ]
+
+
+def build_classification_annotation_specs(dataset: 'DatamintBaseDataset') -> list[CategoryAnnotationSpec]:
+    """Build one CategoryAnnotationSpec per classification identifier in the dataset."""
+    groups: dict[str, list[str]] = {}
+    for identifier, value in dataset.image_categories_set:
+        groups.setdefault(identifier, []).append(value)
+    return [
+        CategoryAnnotationSpec(
+            type=AnnotationType.CATEGORY,
+            scope='image',
+            identifier=ident,
+            required=True,
+            values=sorted(vals),
+        )
+        for ident, vals in groups.items()
+    ]
+
+
+def build_annotation_specs_for_task(task_type: 'TaskType | str | None',
+                                    dataset: 'DatamintBaseDataset') -> list[AnnotationSpec] | None:
+    """Derive annotation specs from a dataset given a (possibly unknown) task type.
+
+    Returns ``None`` -- rather than raising -- when ``task_type`` is missing or
+    unrecognized.
+    """
+    if task_type is None:
+        return None
+    if isinstance(task_type, str):
+        try:
+            task_type = TaskType(task_type)
+        except ValueError:
+            return None
+
+    if task_type in (TaskType.IMAGE_CLASSIFICATION, TaskType.MULTILABEL_IMAGE_CLASSIFICATION):
+        return build_classification_annotation_specs(dataset)
+    if task_type == TaskType.IMAGE_SEGMENTATION:
+        return build_segmentation_annotation_specs(dataset, scope='image')
+    if task_type == TaskType.VOLUME_SEGMENTATION:
+        return build_segmentation_annotation_specs(dataset, scope='volume')
+    if task_type in (TaskType.OBJECT_DETECTION, TaskType.INSTANCE_SEGMENTATION):
+        return build_detection_annotation_specs(dataset)
+
+    return None

docs/source/client_api_content.rst

@@ -534,6 +534,45 @@ Models are also created automatically when you pass ``--ai-model`` to
 :doc:`command_line_tools` (``datamint upload``) with a name that doesn't
 exist yet.
 
+Log a model manually
++++++++++++++++++++++
+
+Models trained through a Datamint :mod:`~datamint.lightning.trainers` are
+logged automatically, annotation specs included. If you trained a model
+yourself (outside a trainer) and want to register it, ``api.models.log_model()``
+opens its own MLflow run and does the registration for you:
+
+.. code-block:: python
+
+    model = api.models.log_model(
+        my_trained_model,
+        project="my-project",
+        model_name="my-model",
+    )
+
+To also attach annotation specs (what the model predicts, which
+segmentation labels, box classes, or categories) without building
+``AnnotationSpec`` objects by hand, pass the dataset you trained on. Specs are
+derived from its actual labels, dispatching on ``my_trained_model.task_type``:
+
+.. code-block:: python
+
+    from datamint.dataset import ImageDataset
+
+    dataset = ImageDataset(project="my-project", return_boxes=True)
+
+    model = api.models.log_model(
+        my_trained_model,
+        project="my-project",
+        model_name="my-model",
+        dataset=dataset,
+    )
+
+Passing ``annotation_specs`` explicitly always takes precedence over anything
+derived from ``dataset``. Omitting ``dataset`` altogether logs the model with
+no annotation specs, same as calling
+:func:`~datamint.mlflow.flavors.datamint_flavor.log_model` directly.
+
 Inspect versions and metrics
 ++++++++++++++++++++++++++++