Feat/mlflow inference

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:

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:

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,23 +57,96 @@ 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,
             'to': date_to.isoformat() if date_to is not None else None,
             '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],

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.

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]: