From 15107c340a104d18082c5f606c8fdb38aa9b2d15 Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Fri, 21 Aug 2026 10:45:46 -0300 Subject: [PATCH] add log model --- datamint/api/endpoints/models_api.py | 96 +++++++++++++++++++ .../trainers/classification_trainer.py | 16 +--- .../lightning/trainers/detection_trainer.py | 7 +- .../segmentation_modules/unetrpp.py | 3 + .../trainers/segmentation_trainer.py | 7 +- .../trainers/specialized/nnunet/trainer.py | 7 +- .../lightning/trainers/vol_seg_trainer.py | 7 +- datamint/mlflow/flavors/annotation_specs.py | 72 ++++++++++++++ docs/source/client_api_content.rst | 39 ++++++++ 9 files changed, 220 insertions(+), 34 deletions(-) create mode 100644 datamint/mlflow/flavors/annotation_specs.py diff --git a/datamint/api/endpoints/models_api.py b/datamint/api/endpoints/models_api.py index 872dd012..ee39d1ee 100644 --- a/datamint/api/endpoints/models_api.py +++ b/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() diff --git a/datamint/lightning/trainers/classification_trainer.py b/datamint/lightning/trainers/classification_trainer.py index 54a04472..c968ce2c 100644 --- a/datamint/lightning/trainers/classification_trainer.py +++ b/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): diff --git a/datamint/lightning/trainers/detection_trainer.py b/datamint/lightning/trainers/detection_trainer.py index d1d3b0ba..ca70d2a5 100644 --- a/datamint/lightning/trainers/detection_trainer.py +++ b/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) diff --git a/datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py b/datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py index 9ad6d4c1..1eaa86d9 100644 --- a/datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py +++ b/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: diff --git a/datamint/lightning/trainers/segmentation_trainer.py b/datamint/lightning/trainers/segmentation_trainer.py index 0e492662..6afaa4af 100644 --- a/datamint/lightning/trainers/segmentation_trainer.py +++ b/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') diff --git a/datamint/lightning/trainers/specialized/nnunet/trainer.py b/datamint/lightning/trainers/specialized/nnunet/trainer.py index 8c7ab19d..0a0eced6 100644 --- a/datamint/lightning/trainers/specialized/nnunet/trainer.py +++ b/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( diff --git a/datamint/lightning/trainers/vol_seg_trainer.py b/datamint/lightning/trainers/vol_seg_trainer.py index 067d9a49..58d4915f 100644 --- a/datamint/lightning/trainers/vol_seg_trainer.py +++ b/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') diff --git a/datamint/mlflow/flavors/annotation_specs.py b/datamint/mlflow/flavors/annotation_specs.py new file mode 100644 index 00000000..419db84b --- /dev/null +++ b/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 diff --git a/docs/source/client_api_content.rst b/docs/source/client_api_content.rst index 35d7a0d5..9cd7f167 100644 --- a/docs/source/client_api_content.rst +++ b/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 ++++++++++++++++++++++++++++