diff --git a/datamint/api/base_api.py b/datamint/api/base_api.py index 0f95b425..78cd3885 100644 --- a/datamint/api/base_api.py +++ b/datamint/api/base_api.py @@ -3,7 +3,6 @@ import httpx from dataclasses import dataclass from datamint.exceptions import DatamintException, ResourceNotFoundError -from datamint.types import ImagingData import aiohttp import json from PIL import Image @@ -17,10 +16,10 @@ if TYPE_CHECKING: from datamint.api.client import Api + from datamint.types import ImagingData logger = logging.getLogger(__name__) -# Generic type for entities _PAGE_LIMIT = 5000 @dataclass @@ -380,9 +379,13 @@ def _make_request_with_pagination(self, """ offset = 0 total_fetched = 0 - params = dict(kwargs.get('params', {})) - # Ensure kwargs carries our params reference so mutations below take effect - kwargs['params'] = params + + use_json_pagination = method.upper() == 'POST' and 'json' in kwargs and isinstance(kwargs['json'], dict) + + if not use_json_pagination: + params = dict(kwargs.get('params', {})) + # Ensure kwargs carries our params reference so mutations below take effect + kwargs['params'] = params while True: if limit is not None and total_fetched >= limit: @@ -393,8 +396,12 @@ def _make_request_with_pagination(self, remaining = limit - total_fetched page_limit = min(_PAGE_LIMIT, remaining) - params['offset'] = offset - params['limit'] = page_limit + if use_json_pagination: + kwargs['json']['offset'] = str(offset) + kwargs['json']['limit'] = str(page_limit) + else: + params['offset'] = offset + params['limit'] = page_limit response = self._make_request(method=method, endpoint=endpoint, @@ -447,7 +454,7 @@ def _convert_array_response(self, def convert_format(bytes_array: bytes, mimetype: str | None = None, file_path: str | None = None - ) -> ImagingData | bytes: + ) -> 'ImagingData | bytes': """ Convert the bytes array to the appropriate format based on the mimetype. Args: @@ -465,6 +472,8 @@ def convert_format(bytes_array: bytes, >>> dicom = BaseApi.convert_format(dicom_bytes) """ + import pydicom + if mimetype is None: mimetype, ext = BaseApi._determine_mimetype(bytes_array) if mimetype is None: diff --git a/datamint/api/client.py b/datamint/api/client.py index 5a62bab0..3e230ee6 100644 --- a/datamint/api/client.py +++ b/datamint/api/client.py @@ -65,7 +65,7 @@ def check_connection(self): self.projects.get_list(limit=1) except Exception as e: raise DatamintException("Error connecting to the Datamint API." + - f" Please check your api_key and/or other configurations. {e}") + f" Please check your api_key and/or other configurations.") from e def _get_endpoint(self, name: str): if self._client is None: diff --git a/datamint/api/endpoints/annotations_api.py b/datamint/api/endpoints/annotations_api.py index 0e213ffa..1c9f8862 100644 --- a/datamint/api/endpoints/annotations_api.py +++ b/datamint/api/endpoints/annotations_api.py @@ -4,7 +4,7 @@ import logging from ..entity_base_api import ApiConfig, CreatableEntityApi, DeletableEntityApi from .models_api import ModelsApi -from datamint.entities.annotation import Annotation +from datamint.entities.annotations.annotation import Annotation from datamint.entities.resource import Resource from datamint.api.dto import AnnotationType, CreateAnnotationDto, LineGeometry, BoxGeometry, CoordinateSystem, Geometry import numpy as np @@ -44,10 +44,11 @@ def __init__(self, from .resources_api import ResourcesApi super().__init__(config, Annotation, 'annotations', client) self._models_api = ModelsApi(config, client=client) if models_api is None else models_api - self._resources_api = ResourcesApi(config, client=client, annotations_api=self) if resources_api is None else resources_api + self._resources_api = ResourcesApi( + config, client=client, annotations_api=self) if resources_api is None else resources_api def get_list(self, - resource: str | Resource | None = None, + resource: str | Resource | Sequence[str | Resource] | None = None, annotation_type: AnnotationType | str | None = None, annotator_email: str | None = None, date_from: date | None = None, @@ -56,10 +57,58 @@ def get_list(self, worklist_id: str | None = None, status: Literal['new', 'published'] | None = None, load_ai_segmentations: bool | None = None, - limit: int | None = None - ) -> Sequence[Annotation]: + limit: int | None = None, + group_by_resource: bool = False + ) -> Sequence[Annotation] | Sequence[Sequence[Annotation]]: + """ + Retrieve a list of annotations with optional filtering. + + Args: + resource: The resource unique id(s) or Resource instance(s). Can be a single resource, + a list of resources, or None to retrieve annotations from all resources. + annotation_type: Filter by annotation type (e.g., 'segmentation', 'category'). + annotator_email: Filter by annotator email address. + date_from: Filter annotations created on or after this date. + date_to: Filter annotations created on or before this date. + dataset_id: Filter by dataset unique id. + worklist_id: Filter by annotation worklist unique id. + status: Filter by annotation status ('new' or 'published'). + load_ai_segmentations: Whether to load AI-generated segmentations. + limit: Maximum number of annotations to return. + group_by_resource: If True, return results grouped by resource. + For instance, the first index of the returned list will contain all annotations for the first resource. + + Returns: + Sequence[Annotation] | Sequence[Sequence[Annotation]]: List of annotations, or list of lists if grouped by resource. + + Example: + .. code-block:: python + + # Get all annotations for a single resource + annotations = api.annotations.get_list(resource='resource_id') + + # Get annotations with filters + annotations = api.annotations.get_list( + resource='resource_id', + annotation_type='segmentation', + status='published' + ) + + # Get annotations for multiple resources + annotations = api.annotations.get_list( + resource=['resource_id_1', 'resource_id_2', 'resource_id_3'] + ) + """ + def group_annotations_by_resource(annotations: Sequence[Annotation], + resource_ids: Sequence[str] + ) -> Sequence[Sequence[Annotation]]: + resource_annotations_map = {rid: [] for rid in resource_ids} + for ann in annotations: + resource_annotations_map[ann.resource_id].append(ann) + return [resource_annotations_map[rid] for rid in resource_ids] + + # Build search payload according to POST /annotations/search schema payload = { - 'resource_id': resource.id if isinstance(resource, Resource) else resource, 'annotation_type': annotation_type, 'annotatorEmail': annotator_email, 'from': date_from.isoformat() if date_from is not None else None, @@ -67,12 +116,37 @@ def get_list(self, 'dataset_id': dataset_id, 'annotation_worklist_id': worklist_id, 'status': status, - 'load_ai_segmentations': load_ai_segmentations + 'load_ai_segmentations': load_ai_segmentations, } - # remove nones + if isinstance(resource, (str, Resource)): + resource_id = self._entid(resource) + payload['resource_id'] = resource_id + resource_ids = None + elif resource is not None: + resource_ids = [self._entid(res) for res in resource] + payload['resource_ids'] = resource_ids + else: + resource_ids = None + + # Remove None values from payload payload = {k: v for k, v in payload.items() if v is not None} - return super().get_list(limit=limit, params=payload) + + items_gen = self._make_request_with_pagination('POST', + f'{self.endpoint_base}/search', + return_field=self.endpoint_base, + limit=limit, + json=payload) + + all_items = [] + for _, items in items_gen: + all_items.extend(items) + + all_annotations = [self._init_entity_obj(**item) for item in all_items] + + if group_by_resource and resource_ids is not None: + return group_annotations_by_resource(all_annotations, resource_ids) + return all_annotations async def _upload_segmentations_async(self, resource: str | Resource, @@ -370,14 +444,21 @@ def create(self, resource: str | Resource, annotation_dto: CreateAnnotationDto | Sequence[CreateAnnotationDto] ) -> str | Sequence[str]: - """Create a new annotation. + """Create one or more annotations for a resource. + + .. warning:: + This is an internal method and should not be used directly by users. + Please use specific annotation creation methods like + :py:meth:`create_image_classification` or :py:meth:`upload_segmentations` instead. Args: - resource: The resource unique id or Resource instance. - annotation_dto: A CreateAnnotationDto instance or a list of such instances. + resource (str | Resource): The resource unique id or Resource instance. + annotation_dto (CreateAnnotationDto | Sequence[CreateAnnotationDto]): + A CreateAnnotationDto instance or a list of such instances to be created. Returns: - The id of the created annotation or a list of ids if multiple annotations were created. + str | Sequence[str]: The id of the created annotation if a single annotation + was provided, or a list of ids if multiple annotations were created. """ annotations = [annotation_dto] if isinstance(annotation_dto, CreateAnnotationDto) else annotation_dto @@ -814,7 +895,7 @@ def create_image_classification(self, model_id=model_id ) - return self.create(resource, annotation_dto) + return self.create(resource, annotation_dto) def add_line_annotation(self, point1: tuple[int, int] | tuple[float, float, float], diff --git a/datamint/api/endpoints/projects_api.py b/datamint/api/endpoints/projects_api.py index 65102c08..f1e2a7e5 100644 --- a/datamint/api/endpoints/projects_api.py +++ b/datamint/api/endpoints/projects_api.py @@ -1,4 +1,4 @@ -from typing import Sequence, Literal, TYPE_CHECKING +from typing import Sequence, Literal, TYPE_CHECKING, overload from ..entity_base_api import ApiConfig, CRUDEntityApi from datamint.entities.project import Project import httpx @@ -38,13 +38,38 @@ def get_project_resources(self, project: Project | str) -> list[Resource]: resources = [self.resources_api._init_entity_obj(**item) for item in resources_data] return resources + + @overload + def create(self, + name: str, + description: str, + resources_ids: list[str] | None = None, + is_active_learning: bool = False, + two_up_display: bool = False, + *, + return_entity: Literal[True] = True + ) -> Project: ... + + @overload + def create(self, + name: str, + description: str, + resources_ids: list[str] | None = None, + is_active_learning: bool = False, + two_up_display: bool = False, + *, + return_entity: Literal[False] + ) -> str: ... + def create(self, name: str, description: str, resources_ids: list[str] | None = None, is_active_learning: bool = False, - two_up_display: bool = False - ) -> str: + two_up_display: bool = False, + *, + return_entity: bool = True + ) -> str | Project: """Create a new project. Args: @@ -53,6 +78,7 @@ def create(self, resources_ids: The list of resource ids to be included in the project. is_active_learning: Whether the project is an active learning project or not. two_up_display: Allow annotators to display multiple resources for annotation. + return_entity: Whether to return the created Project instance or just its ID. Returns: The id of the created project. @@ -72,7 +98,7 @@ def create(self, "require_review": False, 'description': description} - return self._create(project_data) + return self._create(project_data, return_entity=return_entity) def get_all(self, limit: int | None = None) -> Sequence[Project]: """Get all projects. diff --git a/datamint/api/endpoints/resources_api.py b/datamint/api/endpoints/resources_api.py index 8925c6d1..cafe29ce 100644 --- a/datamint/api/endpoints/resources_api.py +++ b/datamint/api/endpoints/resources_api.py @@ -2,7 +2,7 @@ from ..base_api import ApiConfig, BaseApi from ..entity_base_api import CreatableEntityApi, DeletableEntityApi from datamint.entities.resource import Resource -from datamint.entities.annotation import Annotation +from datamint.entities.annotations.annotation import Annotation from datamint.exceptions import DatamintException, ResourceNotFoundError from datamint.api.dto import AnnotationType import httpx @@ -147,7 +147,7 @@ def get_list(self, return super().get_list(limit=limit, params=payload) - def get_annotations(self, resource: str | Resource, + def get_annotations(self, resource: str | Resource, annotation_type: AnnotationType | str | None = None) -> Sequence[Annotation]: """Get annotations for a specific resource. @@ -1008,19 +1008,21 @@ def set_tags(self, resource: The resource object or a list of resources. tags: The tags to set. """ - data = {'tags': tags} + + uniq_tags = set(tags) # remove duplicates + if isinstance(resource, Sequence): resource_ids = [self._entid(res) for res in resource] response = self._make_request('PUT', f'{self.endpoint_base}/tags', json={'resource_ids': resource_ids, - 'tags': tags}) + 'tags': list(uniq_tags)}) else: resource_id = self._entid(resource) response = self._make_entity_request('PUT', resource_id, add_path='tags', - json=data) + json={'tags': list(uniq_tags)}) return response # def get_projects(self, resource: Resource) -> Sequence[Project]: diff --git a/datamint/api/entity_base_api.py b/datamint/api/entity_base_api.py index 0abe026b..4254d48d 100644 --- a/datamint/api/entity_base_api.py +++ b/datamint/api/entity_base_api.py @@ -1,4 +1,4 @@ -from typing import Any, TypeVar, Generic, Type, Sequence +from typing import Any, Literal, TypeVar, Generic, Type, Sequence, AsyncGenerator, overload import logging import httpx from datamint.entities.base_entity import BaseEntity @@ -7,7 +7,6 @@ import asyncio from .base_api import ApiConfig, BaseApi import contextlib -from typing import AsyncGenerator logger = logging.getLogger(__name__) T = TypeVar('T', bound=BaseEntity) @@ -248,7 +247,16 @@ class CreatableEntityApi(EntityBaseApi[T]): This class adds methods to handle creation of new entities. """ - def _create(self, entity_data: dict[str, Any]) -> str | list[str | dict]: + @overload + def _create(self, entity_data: dict[str, Any], + return_entity: Literal[True] = True) -> T | list[T]: ... + + @overload + def _create(self, entity_data: dict[str, Any], + return_entity: Literal[False]) -> str | list: ... + + def _create(self, entity_data: dict[str, Any], + return_entity: bool = False) -> str | T | list: """Create a new entity. Args: @@ -263,14 +271,35 @@ def _create(self, entity_data: dict[str, Any]) -> str | list[str | dict]: response = self._make_request('POST', f'/{self.endpoint_base}', json=entity_data) respdata = response.json() if isinstance(respdata, str): + if return_entity: + return self.get_by_id(respdata) return respdata if isinstance(respdata, list): + if return_entity: + logger.warning("Current implementation is slow when returning entities on bulk create." + " Try ``return_entity=False`` for better performance.") + return [self.get_by_id(item['id']) if isinstance(item, dict) and 'id' in item else self.get_by_id(item) + for item in respdata] return respdata if isinstance(respdata, dict): + if return_entity: + try: + return self._init_entity_obj(**respdata) + except: + logger.debug("Failed to init entity obj on create response. Falling back to get_by_id.") + return self.get_by_id(respdata.get('id')) return respdata.get('id') return respdata - def create(self, *args, **kwargs) -> str | T: + @overload + def create(self, *args, return_entity: Literal[True] = True, **kwargs) -> T: ... + + @overload + def create(self, *args, return_entity: Literal[False], **kwargs) -> str: ... + + def create(self, *args, + return_entity: bool = True, + **kwargs) -> str | T: raise NotImplementedError("Subclasses must implement the create method with their own custom parameters") diff --git a/datamint/configs.py b/datamint/configs.py index 87914460..aa7a21cc 100644 --- a/datamint/configs.py +++ b/datamint/configs.py @@ -1,10 +1,9 @@ import yaml import os import logging -from netrc import netrc from platformdirs import PlatformDirs -from typing import Dict from pathlib import Path +from typing import Any APIURL_KEY = 'default_api_url' APIKEY_KEY = 'api_key' @@ -14,6 +13,10 @@ APIURL_KEY: 'DATAMINT_API_URL' } +DEFAULT_VALUES = { + APIURL_KEY: 'https://api.datamint.io' +} + _LOGGER = logging.getLogger(__name__) DIRS = PlatformDirs(appname='datamintapi') @@ -23,17 +26,17 @@ except Exception as e: _LOGGER.error(f"Could not determine home directory: {e}") DATAMINT_DATA_DIR = None - def get_env_var_name(key: str) -> str: return ENV_VARS[key] -def read_config() -> Dict: + +def read_config() -> dict[str, Any]: if os.path.exists(CONFIG_FILE): with open(CONFIG_FILE, 'r') as configfile: return yaml.safe_load(configfile) - return {} + return DEFAULT_VALUES.copy() def set_value(key: str, diff --git a/datamint/dataset/dataset.py b/datamint/dataset/dataset.py index 40a9c90f..99aec30a 100644 --- a/datamint/dataset/dataset.py +++ b/datamint/dataset/dataset.py @@ -7,7 +7,7 @@ import logging from PIL import Image import albumentations -from datamint.entities.annotation import Annotation +from datamint.entities.annotations.annotation import Annotation from medimgkit.readers import read_array_normalized _LOGGER = logging.getLogger(__name__) diff --git a/datamint/entities/__init__.py b/datamint/entities/__init__.py index 7f655673..9b1af1ee 100644 --- a/datamint/entities/__init__.py +++ b/datamint/entities/__init__.py @@ -1,6 +1,6 @@ """DataMint entities package.""" -from .annotation import Annotation +from .annotations.annotation import Annotation from .base_entity import BaseEntity from .channel import Channel, ChannelResourceData from .project import Project diff --git a/datamint/entities/annotations/__init__.py b/datamint/entities/annotations/__init__.py new file mode 100644 index 00000000..31a57f8e --- /dev/null +++ b/datamint/entities/annotations/__init__.py @@ -0,0 +1,9 @@ +from .image_classification import ImageClassification +from .annotation import Annotation +from datamint.api.dto import AnnotationType # FIXME: move this to this module + +__all__ = [ + "ImageClassification", + "Annotation", + "AnnotationType", +] diff --git a/datamint/entities/annotation.py b/datamint/entities/annotations/annotation.py similarity index 84% rename from datamint/entities/annotation.py rename to datamint/entities/annotations/annotation.py index d09f6677..bc3ac293 100644 --- a/datamint/entities/annotation.py +++ b/datamint/entities/annotations/annotation.py @@ -9,16 +9,17 @@ import logging import os -from .base_entity import BaseEntity, MISSING_FIELD -from .cache_manager import CacheManager +from ..base_entity import BaseEntity, MISSING_FIELD +from ..cache_manager import CacheManager from pydantic import PrivateAttr from datetime import datetime from datamint.api.dto import AnnotationType from datamint.types import ImagingData + if TYPE_CHECKING: from datamint.api.endpoints.annotations_api import AnnotationsApi - from .resource import Resource + from ..resource import Resource logger = logging.getLogger(__name__) @@ -33,7 +34,28 @@ _ANNOTATION_CACHE_KEY = "annotation_data" -class Annotation(BaseEntity): +class AnnotationBase(BaseEntity): + """Minimal base class for creating annotations. + + This class contains only the essential fields needed to create annotations. + Use this for creating specific annotation types like ImageClassification. + """ + + identifier: str + scope: str + annotation_type: AnnotationType + confiability: float = 1.0 + + def __init__(self, **data): + """Initialize the annotation base entity.""" + super().__init__(**data) + + @property + def name(self) -> str: + """Get the annotation name (alias for identifier).""" + return self.identifier + +class Annotation(AnnotationBase): """Pydantic Model representing a DataMint annotation. Attributes: @@ -67,32 +89,31 @@ class Annotation(BaseEntity): values: Optional extra values payload for flexible schemas. """ - id: str + id: str | None = None identifier: str scope: str - frame_index: int | None - annotation_type: AnnotationType - text_value: str | None - numeric_value: float | int | None - units: str | None - geometry: list | dict | None - created_at: str # ISO timestamp string - created_by: str - annotation_worklist_id: str | None - status: str - approved_at: str | None # ISO timestamp string - approved_by: str | None - resource_id: str - associated_file: str | None - deleted: bool - deleted_at: str | None # ISO timestamp string - deleted_by: str | None - created_by_model: str | None - set_name: str | None - resource_filename: str | None - resource_modality: str | None - annotation_worklist_name: str | None - user_info: dict | None + frame_index: int | None = None + text_value: str | None = None + numeric_value: float | int | None = None + units: str | None = None + geometry: list | dict | None = None + created_at: str | None = None # ISO timestamp string + created_by: str | None = None + annotation_worklist_id: str | None = None + status: str | None = None + approved_at: str | None = None # ISO timestamp string + approved_by: str | None = None + resource_id: str | None = None + associated_file: str | None = None + deleted: bool = False + deleted_at: str | None = None # ISO timestamp string + deleted_by: str | None = None + created_by_model: str | None = None + set_name: str | None = None + resource_filename: str | None = None + resource_modality: str | None = None + annotation_worklist_name: str | None = None + user_info: dict | None = None values: list | None = MISSING_FIELD file: str | None = None @@ -190,11 +211,6 @@ def type(self) -> str: """Alias for :attr:`annotation_type`.""" return self.annotation_type - @property - def name(self) -> str: - """Get the annotation name (alias for identifier).""" - return self.identifier - @property def index(self) -> int | None: """Get the frame index (alias for frame_index).""" diff --git a/datamint/entities/annotations/image_classification.py b/datamint/entities/annotations/image_classification.py new file mode 100644 index 00000000..b39b717f --- /dev/null +++ b/datamint/entities/annotations/image_classification.py @@ -0,0 +1,12 @@ +from .annotation import Annotation +from datamint.api.dto import AnnotationType + + +class ImageClassification(Annotation): + def __init__(self, + name: str, + value: str, + confiability: float = 1.0): + super().__init__(identifier=name, text_value=value, scope='image', + confiability=confiability, + annotation_type=AnnotationType.CATEGORY) diff --git a/datamint/entities/base_entity.py b/datamint/entities/base_entity.py index f4eba60a..66671e7f 100644 --- a/datamint/entities/base_entity.py +++ b/datamint/entities/base_entity.py @@ -31,13 +31,24 @@ class BaseEntity(BaseModel): are created through API endpoints. """ - model_config = ConfigDict(extra='allow', arbitrary_types_allowed=True) # Allow extra fields and arbitrary types + model_config = ConfigDict(extra='allow', + arbitrary_types_allowed=True, # Allow extra fields and arbitrary types + ser_json_bytes='base64', + val_json_bytes='base64') _api: 'EntityBaseApi[Self] | EntityBaseApi' = PrivateAttr() + def __init__(self, **data): + super().__init__(**data) + # check attributes for MISSING_FIELD and delete them + for field_name in self.__pydantic_fields__.keys(): + if hasattr(self, field_name) and getattr(self, field_name) == MISSING_FIELD: + delattr(self, field_name) + def asdict(self) -> dict[str, Any]: """Convert the entity to a dictionary, including unknown fields.""" - return self.model_dump(warnings='none') + d = self.model_dump(warnings='none') + return {k: v for k, v in d.items() if v != MISSING_FIELD} def asjson(self) -> str: """Convert the entity to a JSON string, including unknown fields.""" @@ -59,10 +70,13 @@ def model_post_init(self, __context: Any) -> None: if have_to_log: _LOGGER.warning(f"Unknown fields {list(self.__pydantic_extra__.keys())} found in {class_name}") - @staticmethod - def is_attr_missing(value: Any) -> bool: + def is_attr_missing(self, attr_name: str) -> bool: """Check if a value is the MISSING_FIELD sentinel.""" - return value == MISSING_FIELD + if attr_name not in self.__pydantic_fields__.keys(): + raise AttributeError(f"Attribute '{attr_name}' not found in entity of type '{self.__class__.__name__}'") + if not hasattr(self, attr_name): + return True + return getattr(self, attr_name) == MISSING_FIELD # deprecated def _refresh(self) -> Self: """Refresh the entity data from the server. @@ -88,5 +102,16 @@ def _ensure_attr(self, attr_name: str) -> None: Args: attr_name: Name of the attribute to check and ensure """ - if self.is_attr_missing(getattr(self, attr_name)): + if attr_name not in self.__pydantic_fields__.keys(): + raise AttributeError(f"Attribute '{attr_name}' not found in entity of type '{self.__class__.__name__}'") + + if self.is_attr_missing(attr_name): self._refresh() + + def has_missing_attrs(self) -> bool: + """Check if the entity has any attributes that are MISSING_FIELD. + + Returns: + True if any attribute is MISSING_FIELD, False otherwise + """ + return any(self.is_attr_missing(attr_name) for attr_name in self.__pydantic_fields__.keys()) diff --git a/datamint/entities/resource.py b/datamint/entities/resource.py index 78ffa341..4bcbc3b0 100644 --- a/datamint/entities/resource.py +++ b/datamint/entities/resource.py @@ -3,18 +3,23 @@ from datetime import datetime from typing import TYPE_CHECKING, Optional, Any, Sequence import logging +import urllib.parse +import urllib.request from .base_entity import BaseEntity, MISSING_FIELD from .cache_manager import CacheManager from pydantic import PrivateAttr -from datamint.api.dto import AnnotationType import webbrowser -from datamint.types import ImagingData +from pathlib import Path +from datamint.api.base_api import BaseApi if TYPE_CHECKING: from datamint.api.endpoints.resources_api import ResourcesApi from .project import Project - from .annotation import Annotation + from .annotations.annotation import Annotation + from datamint.types import ImagingData + from datamint.api.dto import AnnotationType + logger = logging.getLogger(__name__) @@ -70,7 +75,6 @@ class Resource(BaseEntity): location: str upload_channel: str filename: str - modality: str mimetype: str size: int upload_mechanism: str @@ -80,26 +84,32 @@ class Resource(BaseEntity): created_by: str published: bool deleted: bool - source_filepath: str | None - metadata: dict - projects: list[dict] = MISSING_FIELD - published_on: str | None - published_by: str | None + # metadata: dict[str,Any] = {} + modality: str | None = None + source_filepath: str | None = None + # projects: list[dict[str, Any]] | None = None + published_on: str | None = None + published_by: str | None = None tags: list[str] | None = None - publish_transforms: Optional[Any] = None + # publish_transforms: dict[str, Any] | None = None deleted_at: Optional[str] = None deleted_by: Optional[str] = None instance_uid: Optional[str] = None series_uid: Optional[str] = None study_uid: Optional[str] = None patient_id: Optional[str] = None - segmentations: Optional[Any] = None # TODO: Define proper type when spec available - measurements: Optional[Any] = None # TODO: Define proper type when spec available - categories: Optional[Any] = None # TODO: Define proper type when spec available - user_info: Optional[dict] = None + # segmentations: Optional[Any] = None # TODO: Define proper type when spec available + # measurements: Optional[Any] = None # TODO: Define proper type when spec available + # categories: Optional[Any] = None # TODO: Define proper type when spec available + user_info: dict[str, str | None] = MISSING_FIELD _api: 'ResourcesApi' = PrivateAttr() + def __new__(cls, *args, **kwargs): + if cls is Resource and ('local_filepath' in kwargs or 'raw_data' in kwargs): + return super().__new__(LocalResource) + return super().__new__(cls) + def __init__(self, **data): """Initialize the resource entity.""" super().__init__(**data) @@ -110,7 +120,7 @@ def fetch_file_data( auto_convert: bool = True, save_path: str | None = None, use_cache: bool = False, - ) -> bytes | ImagingData: + ) -> 'bytes | ImagingData': """Get the file data for this resource. This method automatically caches the file data locally. On subsequent @@ -148,10 +158,10 @@ def fetch_file_data( if auto_convert: try: - mimetype, ext = self._api._determine_mimetype(img_data, self) - img_data = self._api.convert_format(img_data, - mimetype=mimetype, - file_path=save_path) + mimetype, ext = BaseApi._determine_mimetype(img_data, self.mimetype) + img_data = BaseApi.convert_format(img_data, + mimetype=mimetype, + file_path=save_path) except Exception as e: logger.error(f"Failed to auto-convert resource {self.id}: {e}") @@ -172,7 +182,7 @@ def _save_into_cache(self, data: bytes) -> None: def fetch_annotations( self, - annotation_type: AnnotationType | str | None = None + annotation_type: 'AnnotationType | str | None' = None ) -> Sequence['Annotation']: """Get annotations associated with this resource.""" @@ -189,7 +199,6 @@ def fetch_annotations( # """ # return self._api.get_projects(self) - def invalidate_cache(self) -> None: """Invalidate cached data for this resource. """ @@ -241,7 +250,7 @@ def __repr__(self) -> str: f"modality='{self.modality}', status='{self.status}', " f"published={self.published})" ) - + @property def url(self) -> str: """Get the URL to access this resource in the DataMint web application.""" @@ -251,3 +260,225 @@ def url(self) -> str: def show(self) -> None: """Open the resource in the default web browser.""" webbrowser.open(self.url) + + @staticmethod + def from_local_file(file_path: str | Path): + """Create a LocalResource instance from a local file path. + + Args: + file_path: Path to the local file + """ + return LocalResource(local_filepath=file_path) + + +class LocalResource(Resource): + """Represents a local resource that hasn't been uploaded to DataMint API yet.""" + + local_filepath: str | None = None + raw_data: bytes | None = None + + def __init__(self, + local_filepath: str | Path | None = None, + raw_data: bytes | None = None, + convert_to_bytes: bool = False, + **kwargs): + """Initialize a local resource from a local file path, URL, or raw data. + + Args: + local_filepath: Path to the local file or URL to an online image + raw_data: Raw bytes of the file data + convert_to_bytes: If True and local_filepath is provided, read file into raw_data + """ + from medimgkit.format_detection import guess_type, DEFAULT_MIME_TYPE + from medimgkit.modality_detector import detect_modality + + if raw_data is None and local_filepath is None: + raise ValueError("Either local_filepath or raw_data must be provided.") + if raw_data is not None and local_filepath is not None: + raise ValueError("Only one of local_filepath or raw_data should be provided.") + + # Check if local_filepath is a URL + if local_filepath is not None: + local_filepath_str = str(local_filepath) + if local_filepath_str.startswith(('http://', 'https://')): + # Download content from URL + logger.debug(f"Downloading resource from URL: {local_filepath_str}") + try: + with urllib.request.urlopen(local_filepath_str) as response: + raw_data = response.read() + # Try to get content-type from response headers + content_type = response.headers.get('Content-Type', '').split(';')[0].strip() + except Exception as e: + raise ValueError(f"Failed to download from URL: {local_filepath_str}") from e + + # Extract filename from URL + parsed_url = urllib.parse.urlparse(local_filepath_str) + url_path = urllib.parse.unquote(parsed_url.path) + filename = Path(url_path).name if url_path else 'downloaded_file' + + # Determine mimetype + mimetype, _ = guess_type(raw_data) + if mimetype is None and content_type: + mimetype = content_type + if mimetype is None: + mimetype = DEFAULT_MIME_TYPE + + default_values = { + 'id': '', + 'resource_uri': '', + 'storage': '', + 'location': local_filepath_str, + 'upload_channel': '', + 'filename': filename, + 'modality': None, + 'mimetype': mimetype, + 'size': len(raw_data), + 'upload_mechanism': '', + 'customer_id': '', + 'status': 'local', + 'created_at': datetime.now().isoformat(), + 'created_by': '', + 'published': False, + 'deleted': False, + 'source_filepath': local_filepath_str, + } + new_kwargs = kwargs.copy() + for key, value in default_values.items(): + new_kwargs.setdefault(key, value) + super(Resource, self).__init__( + local_filepath=None, + raw_data=raw_data, + **new_kwargs + ) + self._cache = None + return + + if convert_to_bytes and local_filepath: + with open(local_filepath, 'rb') as f: + raw_data = f.read() + local_filepath = None + if raw_data is not None: + # import io + if isinstance(raw_data, str): + mimetype, _ = guess_type(raw_data.encode()) + else: + mimetype, _ = guess_type(raw_data) + default_values = { + 'id': '', + 'resource_uri': '', + 'storage': '', + 'location': '', + 'upload_channel': '', + 'filename': 'raw_data', + 'modality': None, + 'mimetype': mimetype if mimetype else DEFAULT_MIME_TYPE, + 'size': len(raw_data), + 'upload_mechanism': '', + 'customer_id': '', + 'status': 'local', + 'created_at': datetime.now().isoformat(), + 'created_by': '', + 'published': False, + 'deleted': False, + 'source_filepath': None, + } + new_kwargs = kwargs.copy() + for key, value in default_values.items(): + new_kwargs.setdefault(key, value) + super().__init__( + local_filepath=None, + raw_data=raw_data, + **new_kwargs + ) + self._cache = None + elif local_filepath is not None: + file_path = Path(local_filepath) + if not file_path.exists(): + raise FileNotFoundError(f"File not found: {file_path}") + + mimetype, _ = guess_type(file_path) + if mimetype is None or mimetype == DEFAULT_MIME_TYPE: + logger.warning(f"Could not determine mimetype for file: {file_path}") + size = file_path.stat().st_size + created_at = datetime.fromtimestamp(file_path.stat().st_ctime).isoformat() + + super().__init__( + id="", + resource_uri="", + storage="", + location=str(file_path), + upload_channel="", + filename=file_path.name, + modality=detect_modality(file_path), + mimetype=mimetype, + size=size, + upload_mechanism="", + customer_id="", + status="local", + created_at=created_at, + created_by="", + published=False, + deleted=False, + source_filepath=str(file_path), + local_filepath=str(file_path), + raw_data=None, + ) + self._cache = None + + def fetch_file_data( + self, *args, + auto_convert: bool = True, + save_path: str | None = None, + **kwargs, + ) -> 'bytes | ImagingData': + """Get the file data for this local resource. + + Args: + auto_convert: If True, automatically converts to appropriate format (pydicom.Dataset, PIL Image, etc.) + save_path: Optional path to save the file locally + Returns: + File data (format depends on auto_convert and file type) + """ + if self.raw_data is not None: + img_data = self.raw_data + local_filepath = None + else: + local_filepath = str(self.local_filepath) + with open(local_filepath, 'rb') as f: + img_data = f.read() + + if save_path: + with open(save_path, 'wb') as f: + f.write(img_data) + + if auto_convert: + try: + mimetype, ext = BaseApi._determine_mimetype(img_data, self.mimetype) + img_data = BaseApi.convert_format(img_data, + mimetype=mimetype, + file_path=local_filepath) + except Exception as e: + logger.error(f"Failed to auto-convert local resource: {e}") + logger.error(e, exc_info=True) + + return img_data + + def __str__(self) -> str: + """String representation of the local resource. + + Returns: + Human-readable string describing the local resource + """ + return f"LocalResource(filepath='{self.local_filepath}', size={self.size_mb}MB)" + + def __repr__(self) -> str: + """Detailed string representation of the local resource. + + Returns: + Detailed string representation for debugging + """ + return ( + f"LocalResource(filepath='{self.local_filepath}', " + f"filename='{self.filename}', modality='{self.modality}', " + f"size={self.size_mb}MB)" + ) diff --git a/datamint/mlflow/__init__.py b/datamint/mlflow/__init__.py index 1c19141c..7ab5e36f 100644 --- a/datamint/mlflow/__init__.py +++ b/datamint/mlflow/__init__.py @@ -4,6 +4,7 @@ from functools import wraps import logging from .env_utils import setup_mlflow_environment, ensure_mlflow_configured +from typing import TYPE_CHECKING _LOGGER = logging.getLogger(__name__) @@ -43,4 +44,18 @@ def _patched_get_tracking_uri(*args, **kwargs): mlflow_utils.get_tracking_uri = _patched_get_tracking_uri -__all__ = ['set_project', 'setup_mlflow_environment', 'ensure_mlflow_configured'] +if TYPE_CHECKING: + from .flavors.model import DatamintModel +else: + import lazy_loader as lazy + + __getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=['flavors.model', 'flavors.datamint_flavor'], + submod_attrs={ + "flavors.model": ["DatamintModel"], + "flavors.datamint_flavor": ["log_model", "load_model"], + }, + ) + +__all__ = ['set_project', 'setup_mlflow_environment', 'ensure_mlflow_configured', 'DatamintModel'] \ No newline at end of file diff --git a/datamint/mlflow/env_utils.py b/datamint/mlflow/env_utils.py index eb58a52b..98542a4e 100644 --- a/datamint/mlflow/env_utils.py +++ b/datamint/mlflow/env_utils.py @@ -15,17 +15,8 @@ def get_datamint_api_url() -> Optional[str]: """Get the Datamint API URL from configuration or environment variables.""" - # First check environment variable - api_url = os.getenv('DATAMINT_API_URL') - if api_url: - return api_url - - # Then check configuration - api_url = configs.get_value(configs.APIURL_KEY) - if api_url: - return api_url - - return None + api_url = configs.get_value(configs.APIURL_KEY, include_envvars=True) # configs checks env vars first + return api_url def get_datamint_api_key() -> Optional[str]: diff --git a/datamint/mlflow/flavors/__init__.py b/datamint/mlflow/flavors/__init__.py new file mode 100644 index 00000000..64fde03a --- /dev/null +++ b/datamint/mlflow/flavors/__init__.py @@ -0,0 +1,17 @@ +""" +Datamint MLflow custom flavor for wrapping PyTorch models with preprocessing. +""" + +from .datamint_flavor import ( + save_model, + log_model, + load_model, + _load_pyfunc, +) + +__all__ = [ + "save_model", + "log_model", + "load_model", + "_load_pyfunc", +] diff --git a/datamint/mlflow/flavors/datamint_flavor.py b/datamint/mlflow/flavors/datamint_flavor.py new file mode 100644 index 00000000..e1ad45c6 --- /dev/null +++ b/datamint/mlflow/flavors/datamint_flavor.py @@ -0,0 +1,156 @@ +import mlflow +from mlflow.models import Model, ModelInputExample, ModelSignature +import datamint +import datamint.mlflow.flavors +from mlflow import pyfunc +from .model import DatamintModel +import logging +from typing import Sequence +from dataclasses import asdict +from packaging.requirements import Requirement + +FLAVOR_NAME = 'datamint' + +_LOGGER = logging.getLogger(__name__) + + +def save_model(datamint_model: DatamintModel, + path, + supported_modes: Sequence[str] | None = None, + data_path=None, + code_paths=None, + infer_code_paths=False, + conda_env=None, + mlflow_model: Model | None = None, + artifacts=None, + signature: ModelSignature | None = None, + input_example: ModelInputExample | None = None, + pip_requirements=None, + extra_pip_requirements=None, + metadata=None, + model_config=None, + example_no_conversion=None, + streamable=None, + **kwargs): + import medimgkit + + if mlflow_model is None: + mlflow_model = Model() + + mlflow_model.add_flavor( + FLAVOR_NAME, + datamint_version=datamint.__version__, + supported_modes=supported_modes or datamint_model.get_supported_modes(), + model_settings=asdict(datamint_model.settings), + linked_models=datamint_model._get_linked_models_uri() + ) + + model_config = model_config or {} + model_config.setdefault('device', 'cuda' if datamint_model.settings.need_gpu else 'cpu') + + def _get_req_name(req): + try: + return Requirement(req).name.lower() + except Exception: + return req.split("==")[0].strip().lower() + + datamint_requirements = ['datamint=={}'.format(datamint.__version__), 'medimgkit=={}'.format(medimgkit.__version__)] + + user_requirements = [] + # Check if requirements are lists (not strings which are also Sequences) + if pip_requirements and isinstance(pip_requirements, Sequence) and not isinstance(pip_requirements, str): + user_requirements.extend(pip_requirements) + if extra_pip_requirements and isinstance(extra_pip_requirements, Sequence) and not isinstance(extra_pip_requirements, str): + user_requirements.extend(extra_pip_requirements) + + user_req_names = {_get_req_name(req) for req in user_requirements} + + missing_requirements = [req for req in datamint_requirements if _get_req_name(req) not in user_req_names] + + if missing_requirements: + if extra_pip_requirements is None: + extra_pip_requirements = missing_requirements + elif isinstance(extra_pip_requirements, Sequence) and not isinstance(extra_pip_requirements, str): + extra_pip_requirements = list(extra_pip_requirements) + missing_requirements + elif pip_requirements and isinstance(pip_requirements, Sequence) and not isinstance(pip_requirements, str): + pip_requirements = list(pip_requirements) + missing_requirements + + + return mlflow.pyfunc.save_model( + path=path, + python_model=datamint_model, + data_path=data_path, + conda_env=conda_env, + mlflow_model=mlflow_model, + # loader_module=None, + artifacts=artifacts, + code_paths=code_paths, + infer_code_paths=infer_code_paths, + signature=signature, + input_example=input_example, + pip_requirements=pip_requirements, + extra_pip_requirements=extra_pip_requirements, + metadata=metadata, + model_config=model_config, + example_no_conversion=example_no_conversion, + streamable=streamable, + **kwargs + ) + + +def log_model( + datamint_model: DatamintModel, + supported_modes: Sequence[str] | None = None, + artifact_path: str = "datamint_model", + data_path=None, + code_paths=None, + infer_code_paths=False, + conda_env=None, + artifacts=None, + registered_model_name: str | None = None, + signature: ModelSignature | None = None, + input_example: ModelInputExample | None = None, + pip_requirements=None, + extra_pip_requirements=None, + metadata=None, + model_config=None, + example_no_conversion=None, + streamable=None, + **kwargs +): + return Model.log( + datamint_model=datamint_model, + supported_modes=supported_modes, + artifact_path=artifact_path, + flavor=datamint.mlflow.flavors.datamint_flavor, + # loader_module=loader_module, + data_path=data_path, + code_paths=code_paths, + artifacts=artifacts, + conda_env=conda_env, + registered_model_name=registered_model_name, + signature=signature, + input_example=input_example, + pip_requirements=pip_requirements, + extra_pip_requirements=extra_pip_requirements, + metadata=metadata, + model_config=model_config, + example_no_conversion=example_no_conversion, + streamable=streamable, + infer_code_paths=infer_code_paths, + **kwargs + ) + + +def load_model(model_uri: str, device: str | None = None) -> DatamintModel: + if device is not None: + model_config = {'device': device} + else: + model_config = None + return mlflow.pyfunc.load_model(model_uri=model_uri, + model_config=model_config + ).unwrap_python_model() + + +def _load_pyfunc(path: str, model_config=None) -> pyfunc.PyFuncModel: + return mlflow.pyfunc.load_model(model_uri=path, model_config=model_config) diff --git a/datamint/mlflow/flavors/model.py b/datamint/mlflow/flavors/model.py new file mode 100644 index 00000000..eb16e6b9 --- /dev/null +++ b/datamint/mlflow/flavors/model.py @@ -0,0 +1,875 @@ +""" +DataMint Model Adapter Module + +This module provides a flexible framework for wrapping ML models to work with DataMint's +annotation system. It supports various prediction modes for different data types and use cases. +""" + +from typing import Any, TypeAlias +from collections.abc import Callable +from abc import ABC, abstractmethod +from enum import Enum +from dataclasses import dataclass +from mlflow.environment_variables import MLFLOW_DEFAULT_PREDICTION_DEVICE +from mlflow.pyfunc import load_model as pyfunc_load_model +from mlflow.pytorch import load_model as pytorch_load_model +from mlflow.pyfunc import PyFuncModel, PythonModel, PythonModelContext +from datamint.entities.annotations import Annotation +from datamint.entities.resource import Resource +import logging +import os + +logger = logging.getLogger(__name__) + +# Type aliases +# AnnotationList: TypeAlias = Sequence[Annotation] +PredictionResult: TypeAlias = list[list[Annotation]] + + +@dataclass +class ModelSettings: + """ + Deployment and inference configuration for DatamintModel. + + These settings are serialized with the model and used by remote MLflow servers + to properly configure the runtime environment. + """ + # Hardware requirements + need_gpu: bool = False + """Whether GPU is required for inference""" + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> 'ModelSettings': + """Create config from dictionary, raising error on unknown keys.""" + valid_fields = {f.name for f in cls.__dataclass_fields__.values()} + invalid_fields = set(data.keys()) - valid_fields + if invalid_fields: + raise ValueError(f"Invalid fields for ModelSettings: {', '.join(sorted(invalid_fields))}") + return cls(**data) + + +class PredictionMode(str, Enum): + """ + Enumeration of supported prediction modes. + + Each mode corresponds to a specific method signature in DatamintModel. + """ + # Standard modes + DEFAULT = 'default' # Default: process entire resource as-is + + # Simple modes + IMAGE = 'image' # Process single 2d image resource + + # Video/temporal modes + FRAME = 'frame' # Extract and process specific frame + FRAME_RANGE = 'frame_range' # Process contiguous frame range + ALL_FRAMES = 'all_frames' # Process all frames independently + TEMPORAL_SEQUENCE = 'temporal_sequence' # Process with temporal context window + + # 3D volume modes + SLICE = 'slice' # Extract and process specific slice + SLICE_RANGE = 'slice_range' # Process contiguous slice range + PRIMARY_SLICE = 'primary_slice' # Process center/primary slice + # MULTI_PLANE = 'multi_plane' # Process multiple anatomical planes + VOLUME = 'volume' # Process entire 3D volume + + # Spatial modes + # ROI = 'roi' # Process single region of interest + # MULTI_ROI = 'multi_roi' # Process multiple regions + # TILE = 'tile' # Split into tiles (whole slide imaging) + # PATCH = 'patch' # Extract patches around points + + # Advanced modes + INTERACTIVE = 'interactive' # With user prompts (SAM-like) + FEW_SHOT = 'few_shot' # With context examples + # MULTI_VIEW = 'multi_view' # Multiple views of same subject + + +class DatamintModel(ABC, PythonModel): + """ + Abstract adapter class for wrapping models to produce Datamint annotations. + + This class provides a flexible framework for integrating ML models with DataMint. + The main `predict()` method routes requests to specific handlers based on the + prediction mode, allowing users to implement only the modes they need. + + Quick Start: + ----------- + ```python + class MyModel(DatamintModel): + def __init__(self): + super().__init__( + mlflow_models_uri={'model': 'models:/MyModel/latest'}, + config=ModelSettings(need_gpu=True) + ) + + def predict_default(self, model_input, **kwargs): + # Access the device for your computation + device = self.inference_device # Reads from MLFLOW_DEFAULT_PREDICTION_DEVICE or defaults to 'cpu' + model = self.mlflow_models['model'].get_raw_model().to(device) + # ... process and return annotations + return predictions + ``` + + Prediction Modes: + ---------------- + Users can request different prediction modes via params['mode']: + + **Default**: Default processing + ```python + model.predict(resources) # or params={'mode': 'default'} + ``` + + **Video Frame**: Extract specific frame + ```python + model.predict(videos, params={'mode': 'frame', 'frame_index': 42}) + ``` + + **3D Slice**: Extract specific slice + ```python + model.predict(volumes, params={'mode': 'slice', 'slice_index': 50, 'axis': 'axial'}) + ``` + + **Interactive**: With prompts + ```python + model.predict(images, params={'mode': 'interactive', 'prompt': {'points': [[x, y]], 'labels': [1]}}) + ``` + + Common Parameters: + ----------------- + - `confidence_threshold` (float): Filter predictions by confidence score + - `batch_size` (int): Batch size for processing + - `render_annotation` (bool): Return annotated images instead of annotations + + Device Configuration: + -------------------- + The device for computation is automatically configured from the + `MLFLOW_DEFAULT_PREDICTION_DEVICE` environment variable. Access it via `self.inference_device`. + Defaults to 'cpu' if not set. + + Implementation Guide: + -------------------- + 1. Implement `predict_default()` - this is required and serves as fallback + 2. Optionally implement specific modes your model supports + 3. Override `_render_annotations()` if you want to support visualization + 4. Use `self.mlflow_models` to access loaded MLflow models + 5. Configure deployment settings via `ModelSettings` + + See individual method docstrings for detailed parameter specifications. + """ + + LINKED_MODELS_DIR = "linked_models" + + def __init__(self, + settings: ModelSettings | dict[str, Any] | None = None, + mlflow_torch_models_uri: dict[str, str] | None = None, + mlflow_models_uri: dict[str, str] | None = None, + ) -> None: + """ + Initialize the DatamintModel adapter. + + Args: + config: ModelSettings instance or dict with deployment settings. + Example: {'need_gpu': True} + mlflow_torch_models_uri: Dictionary mapping model names to PyTorch model URIs. + Example: {'backbone': 'models:/MyClassifier/2'} + These models will be lazy-loaded and accessible via ``self.mlflow_torch_models_uri['backbone']`` + mlflow_models_uri: Dictionary mapping model names to MLflow URIs. + Example: {'detector': 'models:/MyDetector/1', + 'classifier': 'models:/MyClassifier/latest'} + These models will be lazy-loaded and accessible via ``self.mlflow_models['detector']`` + + """ + super().__init__() + self.mlflow_models_uri = (mlflow_models_uri or {}).copy() + self.mlflow_torch_models_uri = (mlflow_torch_models_uri or {}).copy() + + # Handle settings - convert dict to ModelSettings if needed + if isinstance(settings, dict): + self.settings = ModelSettings.from_dict(settings) + elif isinstance(settings, ModelSettings): + self.settings = settings + else: + self.settings = ModelSettings() + + self._supported_modes_cache = None + + def load_context(self, context: PythonModelContext): + """ + Called by MLflow when loading the model. + + Override this if you need custom loading logic. + """ + self._inference_device = self._load_inference_device(context=context) + self._mlflow_models = self._load_mlflow_models() + self._mlflow_torch_models = self._load_mlflow_torch_models() + # model_config = context.model_config + + def _get_linked_models_uri(self) -> dict[str, Any]: + """Get all linked models (MLflow and PyTorch)""" + linked = {} + linked.update(self.mlflow_models_uri) + linked.update(self.mlflow_torch_models_uri) + return linked + + def _clear_linked_models_cache(self): + """Clear loaded linked models to free memory""" + if hasattr(self, '_mlflow_models'): + del self._mlflow_models + if hasattr(self, '_mlflow_torch_models'): + del self._mlflow_torch_models + + def __getstate__(self): + state = self.__dict__.copy() + + state.pop('_mlflow_models', None) + state.pop('_mlflow_torch_models', None) + + return state + + def __setstate__(self, state): + self.__dict__.update(state) + # avoid possible invalid states after unpickling + self._clear_linked_models_cache() + + def _load_inference_device(self, context: PythonModelContext | None = None) -> str: + """ + Load inference device from model config or environment variable. + """ + import torch + + device = None + if context and context.model_config: + device = context.model_config.get("device", None) + logger.info(f"Model config device: {device}") + if device is None: + env_device = MLFLOW_DEFAULT_PREDICTION_DEVICE.get() + if env_device: + device = env_device + elif torch.cuda.is_available(): + device = 'cuda' + else: + device = 'cpu' + + logger.info(f"Set inference device: {device}") + return device + + @property + def inference_device(self) -> str: + if hasattr(self, '_inference_device') and self._inference_device is not None: + return self._inference_device + env_device = MLFLOW_DEFAULT_PREDICTION_DEVICE.get() + if env_device: + logger.info(f"Inference device not set; getting from environment variable ({env_device})") + return env_device + logger.warning("Inference device not set; defaulting to 'cpu'") + return 'cpu' + + def _load_models_generic(self, uris: dict[str, str], + loader_func: Callable, + **loader_kwargs) -> dict[str, Any]: + """Generic helper to load models from URIs.""" + loaded_models = {} + for name, uri in uris.items(): + model_uri = uri + if os.path.exists(uri): + logger.info(f"Model '{name}' found locally at '{uri}'") + model_uri = os.path.abspath(uri) + elif uri.startswith("models:/"): + local_path = uri.replace("models:/", DatamintModel.LINKED_MODELS_DIR + "/", 1) + if os.path.exists(local_path): + logger.info(f"Model '{name}' found locally at '{local_path}'") + model_uri = os.path.abspath(local_path) + + try: + loaded_models[name] = loader_func(model_uri, **loader_kwargs) + logger.info(f"Loaded model '{name}' from {model_uri}") + except Exception as e: + logger.error(f"Failed to load model '{name}' from {model_uri}: {e}") + raise + return loaded_models + + def _load_mlflow_models(self) -> dict[str, PyFuncModel]: + """Load all MLflow models specified in mlflow_models_uri.""" + return self._load_models_generic( + self.mlflow_models_uri, + pyfunc_load_model, + model_config={'device': self.inference_device} + ) + + def _load_mlflow_torch_models(self) -> dict[str, Any]: + """Load all MLflow PyTorch models specified in mlflow_torch_models_uri.""" + models = self._load_models_generic( + self.mlflow_torch_models_uri, + pytorch_load_model, + device=self.inference_device, + map_location=self.inference_device, + ) + for m in models.values(): + if hasattr(m, 'eval'): + m.eval() + return models + + def get_mlflow_models(self) -> dict[str, PyFuncModel]: + """ + Access loaded MLflow models. + + Returns: + Dictionary mapping model names to PyFuncModel instances. + Use .get_raw_model() to access the underlying model (e.g., torch.nn.Module) + """ + if not hasattr(self, '_mlflow_models'): + logger.warning("Loading MLflow models on first access") + self._mlflow_models = self._load_mlflow_models() + return self._mlflow_models + + def get_mlflow_torch_models(self) -> dict[str, Any]: + """ + Access loaded MLflow PyTorch models. + + Returns: + Dictionary mapping model names to PyTorch model instances. + """ + if not hasattr(self, '_mlflow_torch_models'): + logger.warning("Loading MLflow PyTorch models on first access") + self._mlflow_torch_models = self._load_mlflow_torch_models() + return self._mlflow_torch_models + + # def _preprocess_input(self, + # model_input: list[InferenceResource | Resource | dict[str, Any]], + # params: dict[str, Any]) -> list[Resource]: + # """ + # Preprocess input to convert to list of Resource objects. + + # Args: + # model_input: List of InferenceResource, Resource, or dict + # params: Additional parameters (unused here) + # Returns: + # List of Resource objects + # """ + # resources = [] + # for item in model_input: + # if isinstance(item, Resource): + # resources.append(item) + # elif isinstance(item, InferenceResource): + # resources.append(item.fabricate_resource()) + # elif isinstance(item, dict): + # if 'local_filepath' in item or item.get('id', None) == '': + # logger.debug(f'Creating LocalResource from dict: {item}') + # resources.append(LocalResource(local_filepath=item['local_filepath'])) + # elif 'upload_channel' in item or 'location' in item or 'storage' in item: + # resources.append(Resource(**item)) + # else: + # resources.append(InferenceResource(**item).fabricate_resource()) + # else: + # raise ValueError(f"Unsupported input type: {type(item)}") + # return resources + + def predict(self, + model_input: list[Resource], + params: dict[str, Any] | None = None) -> PredictionResult: + """ + Main prediction entry point. + + Routes to appropriate prediction method based on params['mode']. + DO NOT override this method - implement specific predict_* methods instead. + + Args: + model_input: List of Resource objects to process + params: Optional configuration dictionary with keys: + - mode (str): Prediction mode (default: 'standard') + - confidence_threshold (float): Filter by confidence + - batch_size (int): Batch size for processing + - render_annotation (bool): Return rendered images + - device (str): Computation device + + mode-specific parameters (see individual method docs) + + Returns: + List of annotation lists (one per resource), or rendered outputs + if render_annotation=True + + Raises: + ValueError: If mode is invalid or required parameters are missing + NotImplementedError: If requested mode is not implemented + """ + params = params or {} + # model_input = self._preprocess_input(model_input, params) + + # Parse and validate mode + mode = self._parse_mode(model_input=model_input, params=params) + + # Route to appropriate prediction method + try: + if not self._is_mode_implemented(mode): + if self._is_mode_implemented(PredictionMode.DEFAULT): + logger.info(f"Mode '{mode.value}' not implemented, falling back to default") + mode = PredictionMode.DEFAULT + else: + raise NotImplementedError + logger.debug(f"Routing to '{mode.value}' mode for {len(model_input)} resources") + result = self._route_prediction(model_input, mode, params) + + # Apply common post-processing + result = self._post_process(result, model_input, params) + + return result + + except NotImplementedError: + available = self.get_supported_modes() + raise NotImplementedError( + f"Prediction mode '{mode.value}' is not supported by this model.\n" + f"Supported modes: {', '.join(available)}\n" + f"Implement predict_{mode.value}() to add support for this mode." + ) + + def _parse_mode(self, + params: dict[str, Any], + model_input: list[Resource] | None = None) -> PredictionMode: + """Parse and validate prediction mode from params.""" + mode_str = params.get('mode', PredictionMode.DEFAULT.value) + try: + is_all_image = all(res.mimetype.startswith('image/') for res in model_input) if model_input else False + except Exception: + is_all_image = False + + logger.debug(f"Parsing prediction mode: '{mode_str}' | {is_all_image=}") + + if mode_str == PredictionMode.DEFAULT.value and is_all_image: + mode_str = PredictionMode.IMAGE.value + + try: + return PredictionMode(mode_str) + except ValueError: + valid_modes = [m.value for m in PredictionMode] + raise ValueError( + f"Invalid prediction mode: '{mode_str}'\n" + f"Valid modes: {', '.join(valid_modes)}" + ) + + def _route_prediction(self, + model_input: list[Resource], + mode: PredictionMode, + params: dict[str, Any]) -> PredictionResult: + """Route to the appropriate prediction method based on mode.""" + + # Extract mode-specific parameters and remove from kwargs + mode_params, common_params = self._extract_mode_params(mode, params) + + # Dispatch to appropriate method + method = self._get_method_for_mode(mode) + + # if method is None or not self._is_mode_implemented(mode, method): + # raise NotImplementedError + + # Call with explicit parameters + return method(model_input, **mode_params, **common_params) + + def _extract_mode_params(self, mode: PredictionMode, params: dict[str, Any]) -> tuple[dict, dict]: + """ + Extract mode-specific and common parameters. + + Returns: + Tuple of (mode_specific_params, common_params) + """ + # Define mode-specific parameter mappings + mode_param_keys = { + PredictionMode.FRAME: ['frame_index'], + PredictionMode.FRAME_RANGE: ['start_frame', 'end_frame', 'step'], + PredictionMode.SLICE: ['slice_index', 'axis'], + PredictionMode.SLICE_RANGE: ['start_index', 'end_index', 'axis', 'step'], + PredictionMode.PRIMARY_SLICE: ['axis'], + PredictionMode.INTERACTIVE: ['prompt'], + PredictionMode.FEW_SHOT: ['context_resources', 'k'], + PredictionMode.TEMPORAL_SEQUENCE: ['center_frame', 'window_size'], + PredictionMode.IMAGE: [], + } + + reserved_keys = {'mode', 'confidence_threshold'} + + # Extract parameters + mode_specific = {} + common = {} + + mode_keys = set(mode_param_keys.get(mode, ())) + + for key, value in params.items(): + if key in reserved_keys: + continue # Skip mode itself and post-processing-only params + if key in mode_keys: + mode_specific[key] = value + else: + common[key] = value + + return mode_specific, common + + def _get_method_for_mode(self, mode: PredictionMode): + """Get the method corresponding to the given prediction mode.""" + method_name = f"predict_{mode.value}" + method = getattr(self, method_name, None) + return method + + def get_supported_modes(self) -> list[str]: + """ + Get list of prediction modes supported by this model. + + Returns: + List of mode names (strings) + """ + if self._supported_modes_cache is not None: + return self._supported_modes_cache + + supported = [] + for mode in PredictionMode: + if self._is_mode_implemented(mode): + supported.append(mode.value) + + self._supported_modes_cache = supported + return supported + + def _is_mode_implemented(self, mode: PredictionMode) -> bool: + """Determine whether the given mode has a concrete implementation.""" + method = self._get_method_for_mode(mode) + if method is None: + return False + + # Check if method is from DatamintModel base class (not overridden) + if hasattr(DatamintModel, method.__name__): + self._get_method_for_mode + base_method = getattr(DatamintModel, method.__name__) + # Method is implemented if it's not the same as base class method + return method.__func__ is not base_method + + return True + + def _post_process(self, + predictions: PredictionResult, + resources: list[Resource], + params: dict[str, Any]) -> PredictionResult: + """Apply common post-processing based on params.""" + + # Apply confidence threshold filtering + conf_threshold = params.get('confidence_threshold') + if conf_threshold is not None: + predictions = [ + [ann for ann in pred_list + if getattr(ann, 'confiability', 1.0) >= conf_threshold] + for pred_list in predictions + ] + logger.debug(f"Applied confidence threshold: {conf_threshold}") + + return predictions + + def predict_default(self, + model_input: list[Resource], + **kwargs) -> PredictionResult: + """ + **OPTIONAL**: Default prediction on entire resources. + + This is the default mode and serves as fallback for unimplemented modes. + Override this method to implement default prediction behavior. + If called without being overridden, raises NotImplementedError. + + Args: + model_input: Resources to process + **kwargs: Additional user-defined parameters + + Returns: + List of annotation lists, one per resource + + Example: + ```python + def predict_default(self, model_input, **kwargs): + dataset = MyDataset(model_input) + dataloader = DataLoader(dataset) + model = self.mlflow_models['model'].get_raw_model() + + predictions = [] + for batch in dataloader: + outputs = model(batch) + predictions.extend(self._outputs_to_annotations(outputs)) + + return predictions + ``` + """ + raise NotImplementedError( + "predict_default() must be implemented in your DatamintModel subclass. " + "This is the default fallback mode for prediction." + ) + + # ======================================================================== + # VIDEO/TEMPORAL MODES + # ======================================================================== + + def predict_frame(self, + model_input: list[Resource], + frame_index: int, + **kwargs) -> PredictionResult: + """ + Process specific frame from video resources. + + Args: + model_input: Video resources + frame_index: Index of frame to extract and process (0-based) + + Returns: + Annotations for the specified frame (one list per resource) + + Example: + ```python + # Extract frame 42 from multiple videos + predictions = model.predict( + videos, + params={'mode': 'frame', 'frame_index': 42} + ) + ``` + """ + logger.warning(f"predict_frame not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_frame_range(self, + model_input: list[Resource], + start_frame: int, + end_frame: int, + step: int = 1, + **kwargs) -> PredictionResult: + """ + Process range of frames from video resources. + + Args: + model_input: Video resources + start_frame: Start frame index (inclusive) + end_frame: End frame index (inclusive) + step: Step size between frames (default: 1) + + Returns: + Annotations for frames in range (may be frame-scoped annotations) + + Example: + ```python + # Process frames 0-100, every 10th frame + predictions = model.predict( + videos, + params={'mode': 'frame_range', + 'start_frame': 0, + 'end_frame': 100, + 'step': 10} + ) + ``` + """ + logger.warning(f"predict_frame_range not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_frame_interval(self, + model_input: list[Resource], + interval: int, + start_frame: int = 0, + end_frame: int | None = None, + **kwargs) -> PredictionResult: + """ + Process every nth frame from video resources. + + Args: + model_input: Video resources + interval: Process every nth frame + start_frame: Starting frame index + end_frame: Ending frame index (None = last frame) + + Returns: + Annotations for sampled frames + + Example: + ```python + # Process every 30th frame (1 fps for 30fps video) + predictions = model.predict( + videos, + params={'mode': 'frame_interval', 'interval': 30} + ) + ``` + """ + logger.warning(f"predict_frame_interval not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_all_frames(self, + model_input: list[Resource], + **kwargs) -> PredictionResult: + """ + Process all frames independently. + + Args: + model_input: Video resources + + Returns: + Annotations for all frames (likely frame-scoped) + + Example: + ```python + # Analyze every frame + predictions = model.predict( + videos, + params={'mode': 'all_frames'} + ) + ``` + """ + logger.warning(f"predict_all_frames not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + # ======================================================================== + # 3D VOLUME MODES + # ======================================================================== + + def predict_slice(self, + model_input: list[Resource], + slice_index: int, + axis: str = 'axial', + **kwargs) -> PredictionResult: + """ + Process specific slice from 3D volume. + + Args: + model_input: 3D volume resources (DICOM series, NIfTI, etc.) + slice_index: Index of slice to extract + axis: Anatomical axis ('axial', 'sagittal', 'coronal') + + Returns: + Annotations for the specified slice + + Example: + ```python + # Extract and analyze axial slice 50 + predictions = model.predict( + ct_scans, + params={'mode': 'slice', + 'slice_index': 50, + 'axis': 'axial'} + ) + ``` + """ + logger.warning(f"predict_slice not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_slice_range(self, + model_input: list[Resource], + start_index: int, + end_index: int, + axis: str = 'axial', + step: int = 1, + **kwargs) -> PredictionResult: + """ + Process range of slices from 3D volume. + + Args: + model_input: 3D volume resources + start_index: Start slice index (inclusive) + end_index: End slice index (inclusive) + axis: Anatomical axis + step: Step size between slices + + Returns: + Annotations for slices in range + """ + logger.warning(f"predict_slice_range not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_volume(self, + model_input: list[Resource], + **kwargs) -> PredictionResult: + """ + Process entire 3D volume. + + For true 3D models (not slice-by-slice). + + Args: + model_input: 3D volume resources + + Returns: + 3D annotations for entire volume + """ + logger.warning(f"predict_volume not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + # ======================================================================== + # ADVANCED MODES + # ======================================================================== + + def predict_interactive(self, + model_input: list[Resource], + prompt: dict[str, Any], + **kwargs) -> PredictionResult: + """ + Interactive prediction with user prompts. + + For models like Segment Anything (SAM) that accept user guidance. + + Args: + model_input: Resources to process + prompt: Prompt dictionary with keys: + - 'points': list of [x, y] coordinates + - 'labels': list of labels (1=foreground, 0=background) + - 'boxes': list of [x1, y1, x2, y2] bounding boxes + - 'masks': list of binary mask arrays + + Returns: + Annotations based on prompts + + Example: + ```python + # Segment based on positive and negative points + predictions = model.predict( + images, + params={'mode': 'interactive', + 'prompt': { + 'points': [[100, 150], [200, 250]], + 'labels': [1, 0] # foreground, background + }} + ) + ``` + """ + logger.warning(f"predict_interactive not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_few_shot(self, + model_input: list[Resource], + context_resources: list[Resource], + k: int = 5, + **kwargs) -> PredictionResult: + """ + Few-shot prediction with context examples. + + For models that can adapt based on a few labeled examples. + + Args: + model_input: Resources to annotate + context_resources: Resources with existing annotations to use as examples + k: Number of examples to use (if more are provided) + + Returns: + Annotations informed by context examples + + Example: + ```python + # Predict using similar annotated examples + predictions = model.predict( + new_images, + params={'mode': 'few_shot', + 'context_resources': annotated_examples, + 'k': 3} + ) + ``` + """ + logger.warning(f"predict_few_shot not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) + + def predict_image(self, + model_input: list[Resource], + **kwargs) -> PredictionResult: + """ + Process single 2D image resources. + + Args: + model_input: 2D image resources + + Returns: + Annotations for each image + """ + logger.warning(f"predict_image not implemented, falling back to predict_default") + return self.predict_default(model_input, **kwargs) diff --git a/datamint/mlflow/lightning/callbacks/modelcheckpoint.py b/datamint/mlflow/lightning/callbacks/modelcheckpoint.py index e695c2e2..0a26b1df 100644 --- a/datamint/mlflow/lightning/callbacks/modelcheckpoint.py +++ b/datamint/mlflow/lightning/callbacks/modelcheckpoint.py @@ -4,19 +4,23 @@ from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository from typing import Literal, Any import inspect -import torch from torch import nn import lightning.pytorch as L from datamint.mlflow.models import log_model_metadata, _get_MLFlowLogger from datamint.mlflow.env_utils import ensure_mlflow_configured -import mlflow +import mlflow.models +import mlflow.exceptions +import mlflow.pytorch import logging +import json +import hashlib from lightning.pytorch.loggers import MLFlowLogger _LOGGER = logging.getLogger(__name__) def help_infer_signature(x): + import torch if isinstance(x, torch.Tensor): return x.detach().cpu().numpy() elif isinstance(x, dict): @@ -32,7 +36,7 @@ def help_infer_signature(x): class MLFlowModelCheckpoint(ModelCheckpoint): def __init__(self, *args, register_model_name: str | None = None, - register_model_on: Literal["train", "val", "test", "predict"] | None = None, + register_model_on: Literal["train", "val", "test", "predict"] = 'test', code_paths: list[str] | None = None, log_model_at_end_only: bool = True, additional_metadata: dict[str, Any] | None = None, @@ -43,7 +47,7 @@ def __init__(self, *args, Args: register_model_name (str | None): The name to register the model under in MLFlow. If None, the model will not be registered. - register_model_on (Literal["train", "val", "test", "predict"] | None): The stage at which to register the model. If None, the model will not be registered. + register_model_on (Literal["train", "val", "test", "predict"]): The stage at which to register the model. It registers at the end of the specified stage. code_paths (list[str] | None): List of paths to Python files that should be included in the MLFlow model. log_model_at_end_only (bool): If True, only log the model to MLFlow at the end of the training instead of after every checkpoint save. additional_metadata (dict[str, Any] | None): Additional metadata to log with the model as a JSON file. @@ -64,13 +68,12 @@ def __init__(self, *args, if register_model_name is not None and register_model_on is None: raise ValueError("If you provide a register_model_name, you must also provide a register_model_on.") - if register_model_on is not None and register_model_name is None: - raise ValueError("If you provide a register_model_on, you must also provide a register_model_name.") - if register_model_on not in ["train", "val", "test", "predict", None]: + if register_model_on not in ["train", "val", "test", "predict"]: raise ValueError("register_model_on must be one of train, val, test or predict.") self.register_model_name = register_model_name self.register_model_on = register_model_on + self.registered_model_info = None self.log_model_at_end_only = log_model_at_end_only self._last_model_uri = None self.last_saved_model_info = None @@ -79,6 +82,52 @@ def __init__(self, *args, self.code_paths = code_paths self.additional_metadata = additional_metadata or {} self.extra_pip_requirements = extra_pip_requirements or [] + self._last_registered_state_hash: str = "None" + self._has_been_trained: bool = False + + def _compute_registration_state_hash(self) -> str: + """Compute a hash representing the current model state for registration comparison. + + Returns: + A hash string of the current state, or None if state cannot be computed. + """ + state_dict = { + 'checkpoint_path': str(self._last_checkpoint_saved), + 'global_step': self._last_global_step_saved, + 'signature': str(self._inferred_signature) if self._inferred_signature else None, + 'model_uri': self._last_model_uri, + } + + state_str = json.dumps(state_dict, sort_keys=True) + return hashlib.md5(state_str.encode('utf-8')).hexdigest() + + def _should_register_model(self) -> bool: + """Determine if the model should be registered. + + Returns: + True if the model should be registered, False otherwise. + """ + + if self._last_model_uri is None: + _LOGGER.warning("No model URI available. Cannot register model.") + return False + + # If never registered before, register + if self._last_registered_state_hash is None: + return True + + # If model was retrained, register + if self._has_been_trained: + return True + + # If state changed (signature, checkpoint, etc.), register + current_state_hash = self._compute_registration_state_hash() + if current_state_hash != self._last_registered_state_hash: + _LOGGER.debug("Model state has changed since last registration, will register.") + return True + + _LOGGER.info("Model already registered with same configuration. Skipping registration.") + return False def _infer_params(self, model: nn.Module) -> tuple[dict, ...]: """Extract metadata from the model's forward method signature. @@ -211,12 +260,24 @@ def _remove_checkpoint(self, trainer: L.Trainer, filepath: str) -> None: def register_model(self, trainer=None): """Register the model in MLFlow Model Registry.""" + if not self._should_register_model(): + return self.registered_model_info + # mlflow_client = _get_MLFlowLogger(trainer)._mlflow_client - return mlflow.register_model( + self.registered_model_info = mlflow.register_model( model_uri=self._last_model_uri, name=self.register_model_name, ) + # Update the registered state hash after successful registration + self._last_registered_state_hash = self._compute_registration_state_hash() + self._has_been_trained = False # Reset training flag after registration + + _LOGGER.info(f"Model registered as '{self.register_model_name}' " + f"version {self.registered_model_info.version}") + + return self.registered_model_info + def _update_signature(self, trainer): if self._inferred_signature is None: _LOGGER.warning("No signature found. Cannot update signature.") @@ -257,7 +318,6 @@ def wrapped_forward(x, *args, **kwargs): self._inferred_signature = mlflow.models.infer_signature(model_input=x0, params=infered_params) - # run once and get back to the original forward pl_module.forward = original_forward method = getattr(pl_module, 'forward') @@ -271,6 +331,7 @@ def wrapped_forward(x, *args, **kwargs): pl_module.forward = wrapped_forward def on_train_start(self, trainer, pl_module): + self._has_been_trained = True self.__wrap_forward(pl_module) def on_train_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: @@ -285,7 +346,7 @@ def on_train_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None self._update_signature(trainer) - if self.register_model_on == 'train': + if self.register_model_on == 'train' and self.register_model_name: self.register_model(trainer) def _restore_model_uri(self, trainer: L.Trainer) -> None: @@ -319,20 +380,20 @@ def on_predict_start(self, trainer, pl_module): def on_test_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: super().on_test_end(trainer, pl_module) - if self.register_model_on == 'test': + if self.register_model_on == 'test' and self.register_model_name: self._update_signature(trainer) self.register_model(trainer) def on_predict_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: super().on_predict_end(trainer, pl_module) - if self.register_model_on == 'predict': + if self.register_model_on == 'predict' and self.register_model_name: self._update_signature(trainer) self.register_model(trainer) def on_validation_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: super().on_validation_end(trainer, pl_module) - if self.register_model_on == 'val': + if self.register_model_on == 'val' and self.register_model_name: self._update_signature(trainer) self.register_model(trainer) diff --git a/datamint/mlflow/tracking/datamint_store.py b/datamint/mlflow/tracking/datamint_store.py index f0413f47..82e96bb1 100644 --- a/datamint/mlflow/tracking/datamint_store.py +++ b/datamint/mlflow/tracking/datamint_store.py @@ -1,6 +1,5 @@ from mlflow.store.tracking.rest_store import RestStore from functools import partial -from .fluent import get_active_project_id import json @@ -25,9 +24,10 @@ def __init__(self, store_uri: str, artifact_uri=None, force_valid=True): get_host_creds = partial(get_default_host_creds, store_uri) super().__init__(get_host_creds=get_host_creds) - def create_experiment(self, name, artifact_location=None, tags=None, project_id: str = None) -> str: + def create_experiment(self, name, artifact_location=None, tags=None, project_id: str | None = None) -> str: from mlflow.protos.service_pb2 import CreateExperiment from mlflow.utils.proto_json_utils import message_to_json + from datamint.mlflow.tracking.fluent import get_active_project_id if self.invalid: return super().create_experiment(name, artifact_location, tags) diff --git a/datamint/mlflow/tracking/fluent.py b/datamint/mlflow/tracking/fluent.py index b39b1819..c45d1d75 100644 --- a/datamint/mlflow/tracking/fluent.py +++ b/datamint/mlflow/tracking/fluent.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, TYPE_CHECKING import threading import logging from datamint import Api @@ -7,6 +7,9 @@ from datamint.mlflow.env_vars import EnvVars from datamint.mlflow.env_utils import ensure_mlflow_configured +if TYPE_CHECKING: + from datamint.entities.project import Project + _PROJECT_LOCK = threading.Lock() _LOGGER = logging.getLogger(__name__) @@ -44,30 +47,40 @@ def _find_project_by_name(project_name: str): return project -def set_project(project_name: Optional[str] = None, project_id: Optional[str] = None): - from mlflow.exceptions import MlflowException +def _get_project_by_name_or_id(project_name_or_id: str) -> 'Project': + dt_client = Api(check_connection=False) + # If length >= 32, likely an ID + if len(project_name_or_id) >= 32 and ' ' not in project_name_or_id: + # Try to get by ID first + project = dt_client.projects.get_by_id(project_name_or_id) + if project is not None: + return project + project = dt_client.projects.get_by_name(project_name_or_id) + if project is None: + raise DatamintException(f"Project '{project_name_or_id}' does not exist.") + return project + + +def set_project(project: 'Project | str'): + """ + Set the active project for the current session. + + Args: + project: The Project instance or project name/ID to set as active. + """ global _ACTIVE_PROJECT_ID # Ensure MLflow is properly configured before proceeding ensure_mlflow_configured() - if project_name is None and project_id is None: - raise MlflowException("You must specify either a project name or a project id") - - if project_name is not None and project_id is not None: - raise MlflowException("You cannot specify both a project name and a project id") - with _PROJECT_LOCK: - dt_client = Api(check_connection=False) - if project_id is None: - project = dt_client.projects.get_by_name(project_name) - if project is None: - raise DatamintException(f"Project with name '{project_name}' does not exist.") + if isinstance(project, str): + project_id = None + project = _get_project_by_name_or_id(project) project_id = project.id else: - project = dt_client.projects.get_by_id(project_id) - if project is None: - raise DatamintException(f"Project with id '{project_id}' does not exist.") + # It's a Project entity + project_id = project.id _ACTIVE_PROJECT_ID = project_id diff --git a/notebooks/use_cases/fracatlas_classification.ipynb b/notebooks/use_cases/fracatlas_classification.ipynb index 6550ee69..4ac2fa48 100644 --- a/notebooks/use_cases/fracatlas_classification.ipynb +++ b/notebooks/use_cases/fracatlas_classification.ipynb @@ -7,7 +7,25 @@ "source": [ "# Fracture Classification with Datamint and FracAtlas Dataset\n", "\n", - "Train a binary classification model to detect fractures in musculoskeletal radiographs using the FracAtlas dataset." + "This notebook demonstrates how to build an end-to-end binary classification pipeline using **Datamint** and the **FracAtlas** dataset. You will learn how to:\n", + "\n", + "1. **Set up a Datamint project** for managing medical imaging data\n", + "2. **Download and upload** the FracAtlas dataset to Datamint\n", + "3. **Create annotations** for classification tasks\n", + "4. **Build a PyTorch Dataset** that integrates with Datamint\n", + "5. **Train a ResNet-18 model** using PyTorch Lightning\n", + "6. **Track experiments** with MLflow integration\n", + "7. **Deploy the model** for inference using Datamint's model serving\n", + "\n", + "## Required Dependencies\n", + "\n", + "```bash\n", + "pip install datamint\n", + "```\n", + "\n", + "## Dataset Overview\n", + "\n", + "The **FracAtlas** dataset contains 4,083 musculoskeletal radiographs annotated for fracture detection, localization, and segmentation. In this notebook, we focus on **binary classification**: determining whether an X-ray image shows a fracture or not." ] }, { @@ -20,7 +38,6 @@ "from datamint import Api\n", "\n", "PROJECT_NAME = \"FracAtlas\"\n", - "\n", "api = Api()" ] }, @@ -29,7 +46,13 @@ "id": "4450f467", "metadata": {}, "source": [ - "## Setup: Create Project and Upload Dataset" + "## 1. Setup: Create Project and Upload Dataset\n", + "\n", + "In this section, we will:\n", + "- Create a new Datamint project (or retrieve an existing one)\n", + "- Download the FracAtlas dataset from Figshare\n", + "- Upload images to Datamint with appropriate tags\n", + "- Create classification annotations for each image" ] }, { @@ -42,12 +65,12 @@ "from datamint.mlflow import set_project\n", "\n", "proj = api.projects.get_by_name(PROJECT_NAME)\n", - "if not proj:\n", + "if proj is None:\n", " print(f\"Creating project '{PROJECT_NAME}'\")\n", " proj = api.projects.create(name=PROJECT_NAME,\n", - " description=\"Project to train a segmentation model on FracAtlas dataset\")\n", + " description=\"Project to train a binary classification model on FracAtlas dataset\")\n", " \n", - "set_project(PROJECT_NAME)" + "set_project(PROJECT_NAME) # important for proper experiment tracking" ] }, { @@ -55,11 +78,17 @@ "id": "cdd506aa", "metadata": {}, "source": [ - "### Download FracAtlas Dataset\n", + "### 1.1 Download FracAtlas Dataset\n", + "\n", + "The FracAtlas dataset is publicly available on Figshare. We'll download and extract it programmatically.\n", + "\n", + "**Dataset Source:** [Figshare Repository](https://doi.org/10.6084/m9.figshare.22363012)\n", "\n", - "Dataset source: [Figshare](https://doi.org/10.6084/m9.figshare.22363012)\n", + "**Citation:** \n", + "> Abedeen, I., et al. (2023). FracAtlas: A Dataset for Fracture Classification, Localization and Segmentation of Musculoskeletal Radiographs. Scientific Data, 10(1). doi:10.1038/s41597-023-02432-4\n", "\n", - "**Citation:** Abedeen, I., et al. (2023). FracAtlas: A Dataset for Fracture Classification, Localization and Segmentation of Musculoskeletal Radiographs. Scientific Data, 10(1). doi:10.1038/s41597-023-02432-4" + "> [!Note]\n", + "> The download may take a few minutes depending on your internet connection (~1.2 GB compressed).\n" ] }, { @@ -98,22 +127,25 @@ "id": "ff63feee", "metadata": {}, "source": [ - "The dataset is structured as follows:\n", + "### 1.2 Dataset Structure\n", "\n", - "```bash\n", + "The extracted dataset has the following structure:\n", + "\n", + "```\n", "FracAtlas/\n", "├── images/\n", - "│ ├── Fractured/\n", + "│ ├── Fractured/ # 717 images with visible fractures\n", "│ │ ├── IMG0000110.jpg\n", "│ │ └── ...\n", - "│ └── Non_fractured/\n", + "│ └── Non_fractured/ # 3,366 images without fractures\n", "│ ├── IMG0002341.jpg\n", "│ └── ...\n", - "├── Utilities/\n", "└── ...\n", "```\n", "\n", - "We are going to use the `images` folder for our binary classification task." + "For this binary classification task, we'll use only the `images` folder, treating:\n", + "- `Fractured/` → Positive class (label: `has_fracture: yes`)\n", + "- `Non_fractured/` → Negative class (label: `has_fracture: no`)" ] }, { @@ -143,7 +175,10 @@ "metadata": {}, "outputs": [], "source": [ - "# Upload non-fractured images with tags for helping us later in the annotation creation\n", + "# Upload non-fractured images to Datamint\n", + "# Tags are optional and arbitrary, but help organize and filter resources later:\n", + "# 'fracatlas' - identifies dataset source\n", + "# 'non-fractured' - identifies class for annotation creation\n", "new_resources_list = api.resources.upload_resources(non_fractured_images_paths,\n", " tags=['fracatlas', 'non-fractured'],\n", " publish_to=proj, # associate the resources to the project\n", @@ -157,17 +192,27 @@ "metadata": {}, "outputs": [], "source": [ - "# Upload fractured images with tags for helping us later in the annotation creation\n", + "# Upload fractured images to Datamint\n", "new_resources_list = api.resources.upload_resources(fractured_images_paths,\n", " tags=['fracatlas', 'fractured'],\n", " publish_to=proj, # associate the resources to the project\n", " progress_bar=True)" ] }, + { + "cell_type": "markdown", + "id": "07e27496", + "metadata": {}, + "source": [ + "### 1.3 Create Classification Annotations\n", + "\n", + "Now we'll create structured annotations for each image." + ] + }, { "cell_type": "code", "execution_count": null, - "id": "a4db69be", + "id": "0ddc29f8", "metadata": {}, "outputs": [], "source": [ @@ -179,23 +224,25 @@ "for res in tqdm(nonfrac_resources_list):\n", " api.annotations.create_image_classification(resource=res,\n", " identifier='has_fracture',\n", - " value='no')" + " value='no')\n", + "# Annotate fractured images with 'has_fracture: yes'\n", + "frac_resources_list = api.resources.get_list(project_name=PROJECT_NAME,\n", + " tags=['fractured'])\n", + "for res in tqdm(frac_resources_list):\n", + " api.annotations.create_image_classification(resource=res,\n", + " identifier='has_fracture',\n", + " value='yes')" ] }, { "cell_type": "code", "execution_count": null, - "id": "0ddc29f8", + "id": "fdb2758f", "metadata": {}, "outputs": [], "source": [ - "# Annotate fractured images with 'has_fracture: yes'\n", - "frac_resources_list = api.resources.get_list(project_name=PROJECT_NAME,\n", - " tags=['fractured'])\n", - "for res in tqdm(frac_resources_list):\n", - " api.annotations.create_image_classification(resource=res,\n", - " identifier='has_fracture',\n", - " value='yes')" + "# Inspect a sample resource to verify the upload\n", + "frac_resources_list[0]" ] }, { @@ -205,7 +252,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Verify annotation was created successfully\n", + "# Verify that annotations were created correctly\n", + "# The annotation should show identifier='has_fracture' and value='yes' or 'no'\n", "api.annotations.get_list(resource=frac_resources_list[0])[0].asdict()" ] }, @@ -214,9 +262,17 @@ "id": "950ae0f4", "metadata": {}, "source": [ - "### Create Train/Val/Test Splits\n", + "### 1.4 Create Train/Validation/Test Splits\n", "\n", - "Split the dataset into 80% training, 10% validation, and 10% testing with class balance maintained across splits." + "To properly evaluate our model, we split the dataset into three non-overlapping subsets:\n", + "\n", + "| Split | Percentage | Purpose |\n", + "|-------|------------|---------|\n", + "| Train | 80% | Model training |\n", + "| Validation | 10% | Hyperparameter tuning, early stopping |\n", + "| Test | 10% | Final model evaluation |\n", + "\n", + "We use a fixed random seed to ensure reproducibility across runs." ] }, { @@ -254,7 +310,10 @@ "id": "c3ab5bad", "metadata": {}, "source": [ - "Tag resources to identify their split assignment." + "We use Datamint tags to persist the split assignments. This ensures that:\n", + "- Splits are consistent across training sessions\n", + "- Multiple team members can access the same splits\n", + "- The split information is stored alongside the data" ] }, { @@ -277,6 +336,7 @@ "outputs": [], "source": [ "# Verify split distribution and class balance\n", + "# A well-balanced split should have similar fractured ratios across all sets\n", "train_resources = api.resources.get_list(tags=['split:train'])\n", "test_resources = api.resources.get_list(tags=['split:test'])\n", "val_resources = api.resources.get_list(tags=['split:val'])\n", @@ -289,8 +349,7 @@ "\n", "print('Training set: total={}, fractured ratio={:.0%}'.format(total_train, train_fractured_ratio))\n", "print('Validation set: total={}, fractured ratio={:.0%}'.format(total_val, val_fractured_ratio))\n", - "print('Test set: total={}, fractured ratio={:.0%}'.format(total_test, test_fractured_ratio))\n", - "\n" + "print('Test set: total={}, fractured ratio={:.0%}'.format(total_test, test_fractured_ratio))" ] }, { @@ -298,14 +357,24 @@ "id": "9f18e7a8", "metadata": {}, "source": [ - "## Dataset Preparation\n", + "## 2. Dataset Preparation\n", + "\n", + "In this section, we'll create a PyTorch-compatible Dataset class that:\n", + "- Fetches images from Datamint on-demand\n", + "- Applies data augmentation during training\n", + "- Extracts classification labels from annotations\n", "\n", - "Define data transforms and create a PyTorch Dataset class for loading images and annotations from Datamint." + "### 2.1 Define Data Transforms\n", + "\n", + "We use [Albumentations](https://albumentations.ai/) for image augmentation, which provides:\n", + "- Fast, optimized transformations\n", + "- A consistent API for both training and inference\n", + "- Easy integration with PyTorch" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "da4a5dde", "metadata": {}, "outputs": [], @@ -327,55 +396,72 @@ "])" ] }, + { + "cell_type": "markdown", + "id": "52547e56", + "metadata": {}, + "source": [ + "### 2.2 Create Custom PyTorch Dataset\n", + "\n", + "The `FracAtlasDataset` class demonstrates how to integrate Datamint with PyTorch's data loading pipeline. Key features:\n", + "\n", + "- **Lazy loading**: Images are downloaded only when accessed (with optional caching)\n", + "- **Annotation extraction**: Labels are retrieved from Datamint's annotation system\n", + "- **Flexible transforms**: Supports any Albumentations pipeline" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "cb8eb723", + "id": "0e6c8c36", "metadata": {}, "outputs": [], "source": [ "import torch.utils.data\n", "import numpy as np\n", + "from typing import Sequence\n", + "from datamint.entities import Resource\n", + "from collections import defaultdict\n", "\n", "\n", "class FracAtlasDataset(torch.utils.data.Dataset):\n", " \"\"\"Load FracAtlas images and annotations from Datamint for classification.\n", - " \n", + "\n", " Args:\n", " project_name (str): Datamint project name\n", " split (str | None): Filter by split tag ('train', 'val', 'test')\n", " transforms: Albumentations transforms to apply\n", " return_annotations (bool): Include class labels in output\n", " \"\"\"\n", + "\n", " def __init__(self,\n", - " project_name: str,\n", + " resources: Sequence[Resource],\n", " split: str | None = None,\n", " transforms=None,\n", " return_annotations=True,\n", " ):\n", " \"\"\"\n", " Args:\n", - " project_name (str): Name of the Datamint project containing the FracAtlas dataset.\n", + " resources (Sequence[Resource]): List of Datamint Resource objects.\n", " split (str | None): If provided, filters resources by the specified split tag ('train', 'val', 'test').\n", " transforms: Albumentations transforms to apply to the images. Optional.\n", " return_annotations (bool): If True, returns class labels along with images.\n", " \"\"\"\n", - " self.api = Api()\n", - " self.project = self.api.projects.get_by_name(project_name)\n", " self.transforms = transforms\n", " self.return_annotations = return_annotations\n", "\n", - " if not self.project:\n", - " raise ValueError(f\"Project '{project_name}' not found.\")\n", - "\n", - " self.resources = self.project.fetch_resources()\n", + " self.resources = resources\n", " if split:\n", " self.resources = [res for res in self.resources if f'split:{split}' in res.tags]\n", " if return_annotations:\n", - " self.category_annotations = []\n", - " for resource in self.resources:\n", - " annotations = resource.fetch_annotations(annotation_type='category')\n", - " self.category_annotations.append(annotations)\n", + " all_resource_annotations = api.annotations.get_list(resource=self.resources,\n", + " annotation_type='category')\n", + " resource_id_to_annotations = defaultdict(list)\n", + " for ann in all_resource_annotations:\n", + " resource_id_to_annotations[ann.resource_id].append(ann)\n", + "\n", + " self.category_annotations = [resource_id_to_annotations.get(resource.id, [])\n", + " for resource in self.resources]\n", "\n", " def __len__(self):\n", " return len(self.resources)\n", @@ -387,6 +473,9 @@ " image_data = resource.fetch_file_data(auto_convert=True,\n", " use_cache=True) # use_cache=True to avoid re-downloading. By default stored at \"~/.datamint/\"\n", " # image_data is auto converted to a PIL Image (since it is a png image file).\n", + " # If the image was a dicom file, it would be converted to a ``pydicom.dataset.FileDataset`` object.\n", + "\n", + " ### YOUR CUSTOM PREPROCESSING CODE HERE. Example below: ###\n", " image_data = image_data.convert('L') # convert to grayscale\n", " # convert to numpy array float32\n", " image_data = np.array(image_data, dtype=np.float32)\n", @@ -395,37 +484,44 @@ " if self.transforms:\n", " image_data = self.transforms(image=image_data)['image']\n", " # image_data.shape: torch.Size([3, 480, 480])\n", + " ### END OF YOUR CUSTOM PREPROCESSING CODE ###\n", "\n", " if not self.return_annotations:\n", " return image_data\n", - " \n", + "\n", + " ### YOUR CUSTOM ANNOTATION EXTRACTION CODE HERE. Example below: ###\n", " # Extract 'has_fracture' annotation\n", " annotations = self.category_annotations[idx]\n", " for ann in annotations:\n", " if ann.identifier == 'has_fracture':\n", " has_fracture = int(ann.value.lower() == 'yes')\n", " return image_data, has_fracture\n", - " raise ValueError(f\"Annotation 'has_fracture' not found for '{resource.filename}'\")" + " raise ValueError(f\"Annotation 'has_fracture' not found for '{resource.filename}'\")\n", + " ### END OF YOUR CUSTOM ANNOTATION EXTRACTION CODE ###" ] }, { "cell_type": "code", "execution_count": null, - "id": "0e6c8c36", + "id": "ea54a939", "metadata": {}, "outputs": [], "source": [ "# dataloaders\n", "from torch.utils.data import DataLoader\n", "\n", - "train_dataset = FracAtlasDataset(project_name=PROJECT_NAME, split='train', transforms=train_transforms)\n", - "train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4)\n", + "batch_size = 8 # Increase batch size if you have enough GPU memory\n", + "num_workers = 4 # Number of workers for data loading. Increase if you have more CPU cores or set to -1 to use all available cores\n", + "\n", + "all_proj_resources = proj.fetch_resources()\n", + "train_dataset = FracAtlasDataset(resources=all_proj_resources, split='train', transforms=train_transforms)\n", + "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)\n", "\n", - "val_dataset = FracAtlasDataset(project_name=PROJECT_NAME, split='val', transforms=test_transforms)\n", - "val_dataloader = DataLoader(val_dataset, batch_size=8, shuffle=False, num_workers=4)\n", + "val_dataset = FracAtlasDataset(resources=all_proj_resources, split='val', transforms=test_transforms)\n", + "val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n", "\n", - "test_dataset = FracAtlasDataset(project_name=PROJECT_NAME, split='test', transforms=test_transforms)\n", - "test_dataloader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=4)" + "test_dataset = FracAtlasDataset(resources=all_proj_resources, split='test', transforms=test_transforms)\n", + "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)" ] }, { @@ -433,7 +529,9 @@ "id": "80abbae7", "metadata": {}, "source": [ - "## Train the model" + "## 3. Model Training\n", + "\n", + "We fine-tune a pre-trained **ResNet-18** for binary classification." ] }, { @@ -441,8 +539,12 @@ "id": "56844db7", "metadata": {}, "source": [ - "We define our model by extending `lightning.LightningModule` which extends PyTorch `nn.Module`, but handles the dirty work for us.\n", - "More details at https://lightning.ai/docs/pytorch/LTS/common/lightning_module.html" + "### 3.1 Define the Model\n", + "\n", + "We use [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/) to organize our training code. Lightning provides:\n", + "- Automatic GPU/CPU handling\n", + "- Built-in logging and checkpointing\n", + "- Clean separation of training logic" ] }, { @@ -455,6 +557,8 @@ "from torchvision.models import resnet18\n", "import torch\n", "import lightning as L\n", + "from torchmetrics import Accuracy\n", + "\n", "\n", "class FracAtlasClassifier(L.LightningModule):\n", " def __init__(self):\n", @@ -470,18 +574,30 @@ " # Our loss function\n", " self.criterion = torch.nn.CrossEntropyLoss()\n", "\n", + " # Metrics (It must be one metric per split)\n", + " self.test_accuracy = Accuracy(num_classes=2, task='multiclass')\n", + " self.val_accuracy = Accuracy(num_classes=2, task='multiclass')\n", + " self.train_accuracy = Accuracy(num_classes=2, task='multiclass')\n", + "\n", " def forward(self, x):\n", " return self.model(x)\n", "\n", - " def _run_step(self, batch):\n", + " def _run_step(self, batch, split: str | None = None):\n", " \"\"\"Common step for training, validation, and testing.\"\"\"\n", " x, y = batch\n", " y_hat = self(x)\n", " loss = self.criterion(y_hat, y)\n", + " if split == 'train':\n", + " self.train_accuracy.update(y_hat, y)\n", + " elif split == 'val':\n", + " self.val_accuracy.update(y_hat, y)\n", + " elif split == 'test':\n", + " self.test_accuracy.update(y_hat, y)\n", + "\n", " return loss\n", "\n", " def training_step(self, batch, batch_idx):\n", - " loss = self._run_step(batch)\n", + " loss = self._run_step(batch, split='train')\n", " self.log(\"train/loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n", " return loss\n", "\n", @@ -493,9 +609,21 @@ " loss = self._run_step(batch)\n", " self.log(\"test/loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n", " return loss\n", - " \n", + "\n", + " def training_epoch_end(self, outputs):\n", + " self.log(\"train/accuracy\", self.train_accuracy.compute())\n", + " self.train_accuracy.reset()\n", + "\n", + " def validation_epoch_end(self, outputs):\n", + " self.log(\"val/accuracy\", self.val_accuracy.compute())\n", + " self.val_accuracy.reset()\n", + "\n", + " def test_epoch_end(self, outputs):\n", + " self.log(\"test/accuracy\", self.test_accuracy.compute())\n", + " self.test_accuracy.reset()\n", + "\n", " def configure_optimizers(self):\n", - " optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)\n", + " optimizer = torch.optim.Adam(self.parameters(), lr=1e-4) # lr=learning rate\n", " return optimizer\n", "\n", "\n", @@ -513,7 +641,7 @@ "from datamint.mlflow.lightning.callbacks import MLFlowModelCheckpoint\n", "from lightning.pytorch.loggers import MLFlowLogger\n", "\n", - "set_project(PROJECT_NAME) # Ensure the project is set\n", + "set_project(PROJECT_NAME) # Ensure the project is set\n", "\n", "# This callback will do the following:\n", "# - Save the best model based on validation loss.\n", @@ -523,23 +651,21 @@ " mode=\"min\", # Save model when monitored metric decreases\n", " save_top_k=1, # Keep only the best model\n", " filename=\"best\", # Checkpoint filename\n", - " save_weights_only=True, # Save only model weights (not optimizer state)\n", - " register_model_name=PROJECT_NAME, # Name for model registry\n", - " register_model_on='test', # Register model after testing\n", - " # code_paths=['my_custom_model.py'], # Include source code with model\n", - " log_model_at_end_only=True, # Log to MLflow only at the end (faster)\n", + " save_weights_only=True, # Save only model weights (no optimizer state)\n", + " register_model_name=PROJECT_NAME, # *Name for model registry\n", + " register_model_on='test', # Register model only when `trainer.test()` is called\n", + " # code_paths=['my_custom_model.py'], # Include source code with model, if you have defined your model in a separate file\n", ")\n", + "# If you want to register the model manually, you can either run:\n", + "# - `mlflow.pytorch.log_model(model, registered_model_name='MYMODELNAME')` after training;\n", + "# - OR call `checkcb.register_model()` after training.\n", "\n", "# Start Training\n", - "# ==============\n", - "\n", "print(\"🚀 Starting training...\")\n", - "mlflow_logger = MLFlowLogger(experiment_name=PROJECT_NAME)\n", + "mlflow_logger = MLFlowLogger(experiment_name=f'{PROJECT_NAME}_training')\n", "trainer = L.Trainer(\n", - " max_epochs=10, # Number of training epochs\n", + " max_epochs=3, # Number of training epochs\n", " logger=mlflow_logger, # MLflow integration\n", - " enable_model_summary=True, # Show model architecture summary\n", - " enable_progress_bar=True, # Show training progress\n", " callbacks=[checkcb], # Include our checkpoint callback\n", " num_sanity_val_steps=0, # Skip validation sanity check\n", ")\n", @@ -561,13 +687,17 @@ "id": "4bbfbc81", "metadata": {}, "source": [ - "While running, \n", - "- you can check saved model locally with name \"best.ckpt\";\n", - "- And you can check experiment details on the Datamint platform:\n", + "### 3.3 Monitor Training Progress\n", + "\n", + "While training runs, you can:\n", + "- Check the saved model locally (`best.ckpt`)\n", + "- View experiment details on the Datamint platform\n", "\n", - "![image.png](attachment:image.png)\n", + "Run `proj.show()` to open the project dashboard in your browser.\n", "\n", - "![image-2.png](attachment:image-2.png)" + "![Experiment Tracking Dashboard](attachment:image.png)\n", + "\n", + "![Model Metrics Visualization](attachment:image-2.png)" ] }, { @@ -577,6 +707,7 @@ "metadata": {}, "outputs": [], "source": [ + "# Open the Datamint project dashboard to view experiments and metrics\n", "proj.show() # Display project details in Datamint platform" ] }, @@ -587,8 +718,12 @@ "metadata": {}, "outputs": [], "source": [ - "# Start Testing. Important to register the best model\n", - "trainer.test(dataloaders=test_dataloader)" + "# Evaluate on test set and register the best model\n", + "# This step is required to trigger model registration automatically (see register_model_on='test' above)\n", + "trainer.test(dataloaders=test_dataloader)\n", + "\n", + "# you can do the registration manually as well:\n", + "# mlflow.pytorch.log_model(model, registered_model_name=PROJECT_NAME)" ] }, { @@ -596,7 +731,11 @@ "id": "95e977c0", "metadata": {}, "source": [ - "# Predicting" + "## 4. Model Inference\n", + "\n", + "### 4.1 Load and Predict with the Registered Model\n", + "\n", + "We load the trained model from MLflow's Model Registry for prediction." ] }, { @@ -616,10 +755,12 @@ " enable_progress_bar=True,\n", ")\n", "\n", + "print('Loading registered model from MLflow Model Registry...')\n", "registered_model = mlflow.pytorch.load_model(f'models:/{PROJECT_NAME}/latest')\n", + "print('Model loaded successfully!')\n", "\n", "# Set up data module for prediction (same as before, but without annotations)\n", - "test_dataset = FracAtlasDataset(project_name=PROJECT_NAME, split='test',\n", + "test_dataset = FracAtlasDataset(all_proj_resources, split='test',\n", " transforms=test_transforms, return_annotations=False)\n", "test_dataloader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=4)\n", "\n", @@ -632,13 +773,218 @@ "\n", "# Option 2: Load from MLflow Model Registry (commented out)\n", "# registered_model = mlflow.pytorch.load_model(f'models:/{PROJECT_NAME}/latest')\n", - "# preds = pred_trainer.predict(registered_model, datamodule=pred_dm)\n", + "# preds = pred_trainer.predict(registered_model, dataloaders=test_dataloader)\n", "\n", "\n", "print(f\"✅ Predictions completed!\")\n", "print(f\"First batch shape: {preds[0].shape}\")\n", "print(f'First batch, class with max probability: {preds[0].argmax(dim=1)}')" ] + }, + { + "cell_type": "markdown", + "id": "4d7ddddd", + "metadata": {}, + "source": [ + "## 5. Model Deployment\n", + "\n", + "For production use, we need to wrap our model in a **Datamint Model Adapter**. This adapter:\n", + "- Standardizes the input/output format\n", + "- Handles resource loading from Datamint\n", + "- Enables deployment via MLflow Model Serving or Datamint's inference API\n", + "\n", + "### 5.1 Create a Datamint Model Adapter\n", + "\n", + "The `DatamintModel` base class provides a consistent interface for model deployment. Override `predict_image()` to define how your model processes inputs and returns annotations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5654fac1", + "metadata": {}, + "outputs": [], + "source": [ + "from datamint.mlflow.flavors.model import DatamintModel\n", + "from datamint.entities.annotations import ImageClassification\n", + "from datamint.entities import Resource\n", + "import lightning as L\n", + "\n", + "\n", + "class FracAtlasAdapter(DatamintModel):\n", + "\n", + " def __init__(self) -> None:\n", + " # you can link to already saved models: `super().__init__(mlflow_models_uri={'pytorch_model': 'models:/FracAtlas/1'})`\n", + " # Or set them here as attribute, which will be saved together with your adapter: `self.pytorch_model = your_model`\n", + "\n", + " # This line automatically loads the model from MLflow Model Registry, which can be accessed via `self.mlflow_models['pytorch_model']`\n", + " super().__init__(mlflow_torch_models_uri={'pytorch_model': 'models:/FracAtlas/latest'},\n", + " settings={'need_gpu': False})\n", + "\n", + " def predict_image(self,\n", + " model_input: list[Resource],\n", + " **kwargs):\n", + " # Use the same code, dataset and transforms as during testing! No rewrite needed!\n", + " dataset = FracAtlasDataset(model_input, transforms=test_transforms, return_annotations=False)\n", + " dataloader = DataLoader(dataset, batch_size=1, shuffle=False)\n", + " pytorch_model = self.get_mlflow_torch_models()['pytorch_model'] # torch.nn.Module\n", + " pytorch_model.eval()\n", + " # You can use ``L.Trainer`` here,\n", + " # but we will use ``L.Fabric`` (very similar), a more lightweight and flexible version of it.\n", + "\n", + " predictor = L.Fabric() # Way more lightweight and flexible than ``L.Trainer``\n", + " predictor.to_device(self.inference_device)\n", + " pytorch_model = predictor.setup_module(pytorch_model)\n", + " dataloader = predictor.setup_dataloaders(dataloader)\n", + "\n", + " preds = []\n", + " with torch.no_grad():\n", + " for batch in dataloader:\n", + " prob = pytorch_model(batch)\n", + " # prob.shape: torch.Size([1, 2])\n", + " # prob is in logits, let's convert to probabilities\n", + " prob = torch.nn.functional.softmax(prob, dim=1)\n", + " for p in prob:\n", + " max_idx = p.argmax()\n", + " val = 'yes' if max_idx == 1 else 'no'\n", + " annot = ImageClassification(name='has_fracture',\n", + " value=val,\n", + " confiability=float(p[max_idx])\n", + " )\n", + " preds.append([annot])\n", + " return preds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e2d3f957", + "metadata": {}, + "outputs": [], + "source": [ + "from datamint.mlflow.flavors import datamint_flavor\n", + "from datamint.mlflow import set_project\n", + "from mlflow import set_experiment\n", + "import mlflow\n", + "from datamint.entities.resource import LocalResource\n", + "\n", + "set_project('FracAtlas')\n", + "set_experiment('FracAtlas_inference') # Doesn't need to be a new experiment. Do whatever you feel is more organized.\n", + "dtmodel = FracAtlasAdapter()\n", + "\n", + "with mlflow.start_run(run_name=\"adapting_fracatlas_model\"):\n", + " modelinfo = datamint_flavor.log_model(\n", + " dtmodel,\n", + " registered_model_name=\"FracAtlas_adapted\", # You almost always want to register this new adapter-model\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "cea0d758", + "metadata": {}, + "source": [ + "## 6. Testing and Serving\n", + "\n", + "### 6.1 Test the Logged Model Locally\n", + "\n", + "Before deploying, verify that the model loads and predicts correctly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7f29884b", + "metadata": {}, + "outputs": [], + "source": [ + "import mlflow\n", + "\n", + "loaded_model = mlflow.pyfunc.load_model('models:/FracAtlas_adapted/latest')\n", + "loaded_model.predict(all_proj_resources[:2]) # just the first two resources" + ] + }, + { + "cell_type": "markdown", + "id": "1c53d87d", + "metadata": {}, + "source": [ + "> [!TIP]\n", + "> You can retrieve the adapter instance for debugging:\n", + "> ```python\n", + "> my_retrieved_adapter = loaded_model.unwrap_python_model()\n", + "> print(type(my_retrieved_adapter)) # \n", + "> ```" + ] + }, + { + "cell_type": "markdown", + "id": "ddfabd37", + "metadata": {}, + "source": [ + "### 6.2 Serve the Model with MLflow\n", + "\n", + "Serve the model locally via REST API:\n", + "\n", + "```bash\n", + "mlflow models serve -m \"models:/FracAtlas_adapted/latest\" -p 5111 --env-manager virtualenv\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f2467a85", + "metadata": {}, + "outputs": [], + "source": [ + "# Example: Send a prediction request to the model server\n", + "import json\n", + "import requests\n", + "from datamint.entities.resource import LocalResource\n", + "\n", + "example_resource = LocalResource(\n", + " 'https://www.radiologymasterclass.co.uk/images/musculoskeletal-images/trauma/transverse_bone_fracture.jpg'\n", + " # '/My/Local/file.jpeg'\n", + ")\n", + "\n", + "payload = json.dumps(\n", + " {\n", + " \"inputs\": [example_resource.model_dump(mode='json')],\n", + " }\n", + ")\n", + "response = requests.post(\n", + " url=f\"http://localhost:5111/invocations\",\n", + " data=payload,\n", + " headers={\"Content-Type\": \"application/json\"},\n", + ")\n", + "response.json()" + ] + }, + { + "cell_type": "markdown", + "id": "905cea65", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "In this notebook, we demonstrated a complete ML pipeline for medical image classification:\n", + "\n", + "| Step | Description |\n", + "|------|-------------|\n", + "| **Data Management** | Uploaded 4,000+ images to Datamint with tags and annotations |\n", + "| **Dataset Creation** | Built a PyTorch Dataset that integrates with Datamint's API |\n", + "| **Model Training** | Fine-tuned ResNet-18 with PyTorch Lightning |\n", + "| **Experiment Tracking** | Logged metrics and models with MLflow integration |\n", + "| **Model Deployment** | Created a deployable adapter for inference |\n", + "\n", + "### See Also\n", + "\n", + "- [Datamint Documentation](https://sonanceai.github.io/datamint-python-api/)\n", + "- [FracAtlas Paper](https://doi.org/10.1038/s41597-023-02432-4)\n", + "- [PyTorch Lightning Guide](https://lightning.ai/docs/pytorch/stable/)\n", + "- [MLflow Model Registry](https://mlflow.org/docs/latest/model-registry.html)" + ] } ], "metadata": { diff --git a/pyproject.toml b/pyproject.toml index a50bb4cb..e4547ddc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "datamint" description = "A library for interacting with the Datamint API, designed for efficient data management, processing and Deep Learning workflows." -version = "2.4.3" +version = "2.5.0" dynamic = ["dependencies"] requires-python = ">=3.10" readme = "README.md" @@ -37,11 +37,12 @@ Deprecated = ">=1.2.0" platformdirs = "^4.0.0" pandas = ">=2.0.0" matplotlib = "*" -lightning = ">=2.0.0, !=2.5.1, !=2.5.1.post0" -mlflow = "^2.0.0" +lightning = {extras = ['extra'], version = ">=2.0.0, !=2.5.1, !=2.5.1.post0"} +# mlflow = ">=3.0.0" # version 2 has security issues +mlflow = "<3.0.0" albumentations = ">=2.0.0" lazy-loader = ">=0.3.0" -medimgkit = ">=0.7.3" +medimgkit = ">=0.8.0" typing_extensions = ">=4.0.0" pydantic = ">=2.6.4" httpx = "*"