Feat/new dataset class

datamint/__init__.py

@@ -5,20 +5,28 @@
 import importlib.metadata
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
+    # Legacy
     from .dataset.dataset import DatamintDataset as Dataset
-    from .apihandler.api_handler import APIHandler
+
     from .api.client import Api
+    # New modular datasets
+    from .dataset.image_dataset import ImageDataset
+    from .dataset.volume_dataset import VolumeDataset
+
 else:
     import lazy_loader as lazy
 
     __getattr__, __dir__, __all__ = lazy.attach(
         __name__,
-        submodules=['dataset', "dataset.dataset", "apihandler.api_handler"],
+        submodules=['dataset', "dataset.dataset"],
         submod_attrs={
+            # Legacy exports
             "dataset.dataset": ["DatamintDataset"],
             "dataset": ['Dataset'],
-            "apihandler.api_handler": ["APIHandler"],
             "api.client": ["Api"],
+            # New modular dataset classes
+            "dataset.image_dataset": ["ImageDataset"],
+            "dataset.volume_dataset": ["VolumeDataset"],
         },
     )
 

datamint/api/base_api.py

@@ -8,12 +8,10 @@
 import json
 from PIL import Image
 import cv2
-import nibabel as nib
 from io import BytesIO
 import gzip
 import contextlib
 import asyncio
-from medimgkit.format_detection import GZIP_MIME_TYPES, DEFAULT_MIME_TYPE, guess_typez, guess_extension
 from datamint.utils.env import ensure_asyncio_loop
 import os
 
@@ -84,7 +82,6 @@ def __init__(self,
         self._aiohttp_connector: aiohttp.TCPConnector | None = None
         self._aiohttp_session: aiohttp.ClientSession | None = None
         ensure_asyncio_loop()
-
 
     @staticmethod
     def _create_client(config: ApiConfig) -> httpx.Client:
@@ -435,7 +432,6 @@ def _check_errors_response_httpx(self,
                 logger.debug("Unable to set message attribute on exception")
                 pass
 
-            logger.error(f"HTTP error {response.status_code} for {url}: {error_msg}")
             status_code = response.status_code
             if status_code in (400, 404):
                 new_error_msg = error_msg.replace('404 Not Found', '')
@@ -687,6 +683,8 @@ def convert_format(bytes_array: bytes,
 
         """
         import pydicom
+        import nibabel as nib
+        from medimgkit.format_detection import GZIP_MIME_TYPES
 
         if mimetype is None:
             mimetype, ext = BaseApi._determine_mimetype(bytes_array)
@@ -722,6 +720,8 @@ def convert_format(bytes_array: bytes,
                 ndata = nib.Nifti1Image.from_stream(f)
                 ndata.get_fdata()  # force loading before IO is closed
                 return ndata
+        elif mimetype == 'application/x-empty':
+            raise ValueError("Empty file content.")
 
         raise ValueError(f"Unsupported mimetype: {mimetype}")
 
@@ -732,21 +732,19 @@ def _determine_mimetype(content: bytes,
 
         Args:
             content: Raw file content bytes
-            declared_mimetype: Optional MIME type declared by the source
+            declared_mimetype: Optional MIME type declared by the source, used as a fallback if content-based detection fails
 
         Returns:
             Tuple of (inferred_mimetype, file_extension)
         """
+        from medimgkit.format_detection import DEFAULT_MIME_TYPE, guess_typez, guess_extension
         # Determine mimetype from file content
         mimetype_list, ext = guess_typez(content, use_magic=True)
         mimetype = mimetype_list[-1]
 
         # get mimetype from resource info if not detected
         if declared_mimetype is not None:
-            if mimetype is None:
-                mimetype = declared_mimetype
-                ext = guess_extension(mimetype)
-            elif mimetype == DEFAULT_MIME_TYPE:
+            if mimetype is None or mimetype == DEFAULT_MIME_TYPE:
                 mimetype = declared_mimetype
                 ext = guess_extension(mimetype)
 

datamint/api/client.py

@@ -181,3 +181,15 @@ def deploy(self) -> DeployModelApi:
     def inference(self) -> InferenceApi:
         """Access model inference endpoints."""
         return self._get_endpoint('inference', is_mlflow=True)
+
+    def __getstate__(self) -> dict:
+        return {
+            'server_url': self.config.server_url,
+            'api_key': self.config.api_key,
+            'timeout': self.config.timeout,
+            'max_retries': self.config.max_retries,
+            'verify_ssl': self.config.verify_ssl,
+        }
+
+    def __setstate__(self, state: dict) -> None:
+        self.__init__(check_connection=False, **state)

datamint/api/dto/__init__.py

@@ -12,7 +12,5 @@
     "Geometry",
     "BoxGeometry",
     "LineGeometry",
-    "CoordinateSystem"
-    "LineGeometry",
-    "CoordinateSystem"
+    "CoordinateSystem",
 ]

datamint/api/endpoints/annotationsets_api.py

@@ -1,11 +1,49 @@
 from datamint.api.base_api import BaseApi
-import logging
+from typing import TYPE_CHECKING, Any
+from datamint.entities import AnnotationSpec
+from collections.abc import Sequence
 
-_LOGGER = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from datamint.entities import Project
 
 
 class AnnotationSetsApi(BaseApi):
-    def get_segmentation_group(self, annotation_set_id: str) -> dict:
-        """Get the segmentation group for a given annotation set ID."""
-        endpoint = f"/annotationsets/{annotation_set_id}/segmentation-group"
+    ENDPOINT_BASE = "/annotationsets"
+
+    def get_segmentation_group(self, annotation_set: 'str | Project') -> dict:
+        """Get the segmentation group for a given annotation set ID or Project."""
+
+        if isinstance(annotation_set, str):
+            annotation_set_id = annotation_set
+        else:
+            annotation_set_id = annotation_set.worklist_id
+
+        endpoint = f"/{self.ENDPOINT_BASE}/{annotation_set_id}/segmentation-group"
         return self._make_request("GET", endpoint).json()
+
+    def get_annotations_specs(self, annotation_set: 'str | Project') -> Sequence[AnnotationSpec]:
+        """Get the annotations specs for a given annotation set ID or Project."""
+
+        if isinstance(annotation_set, str):
+            annotation_set_id = annotation_set
+        else:
+            annotation_set_id = annotation_set.worklist_id
+
+        result = self._get_by_id(annotation_set_id)
+
+        return [AnnotationSpec(**annspec) for annspec in result['annotations']]
+
+    def _get_by_id(self, annotation_set_id: str) -> dict[str, Any]:
+        """Get an annotation set by its ID.
+
+        Args:
+            annotation_set_id: The ID of the annotation set to retrieve.
+
+        Returns:
+            A dictionary representing the annotation set.
+        """
+        endpoint = f"/{self.ENDPOINT_BASE}/{annotation_set_id}"
+        result = self._make_request("GET", endpoint).json()
+
+        result['annotations'] = [AnnotationSpec(**annspec) for annspec in result['annotations']]
+        return result

datamint/api/endpoints/deploy_model_api.py

@@ -1,4 +1,6 @@
 import httpx
+
+from datamint.exceptions import ResourceNotFoundError
 from ..entity_base_api import EntityBaseApi, ApiConfig
 from datamint.entities.deployjob import DeployJob
 
@@ -16,7 +18,13 @@ def get_by_id(self, entity_id: str) -> DeployJob:
         data = response.json()
         if 'job_id' in data:
             data['id'] = data.pop('job_id')
-        return self._init_entity_obj(**data)
+        self._validate_uuid(data['id'])
+        try:
+            return self._init_entity_obj(**data)
+        except ResourceNotFoundError as e:
+            e.resource_type = 'DeployJob'
+            e.params = {'id': entity_id}
+            raise
 
     def start(self,
               model_name: str,

datamint/api/endpoints/inference_api.py

@@ -75,7 +75,7 @@ def submit(
         model_version: int | None = None,
         model_alias: str | None = None,
         resource_id: str | None = None,
-        resource_ids: list[str] | None = None,
+        # resource_ids: list[str] | None = None,
         file_path: str | None = None,
         file_paths: list[str] | None = None,
         save_results: bool = False,
@@ -106,8 +106,8 @@ def submit(
             save_results=save_results,
             params=params,
         )
-        if resource_ids is not None:
-            payload["resource_ids"] = resource_ids
+        # if resource_ids is not None:
+        #     payload["resource_ids"] = resource_ids
         if file_paths is not None:
             payload["file_paths"] = file_paths
 
@@ -160,49 +160,55 @@ def wait(
         *,
         on_status: Callable[[InferenceJob], None] | None = None,
         poll_interval: float = 2.0,
-        timeout: float | None = None,
-    ) -> InferenceJob:
+        timeout: float | None = 1800,
+    ) -> None:
         """Block until an inference job reaches a terminal state.
 
         First attempts to follow the SSE stream. If the stream is
         unavailable or drops early the method falls back to polling
         ``get_status`` at *poll_interval* seconds.
 
         Args:
-            job: Job ID string or ``InferenceJob`` entity.
+            job: Job ID string or ``InferenceJob`` entity. In-place updates to the provided ``InferenceJob`` are made on every status change.
             on_status: Optional callback invoked with an updated
                 ``InferenceJob`` each time a status update is received.
             poll_interval: Seconds between polls when falling back to
                 polling mode. Default ``2.0``.
             timeout: Maximum seconds to wait. ``None`` means wait
                 indefinitely. Raises ``TimeoutError`` on expiry.
 
-        Returns:
-            The ``InferenceJob`` in its terminal state.
-
         Raises:
             TimeoutError: If *timeout* is set and the job has not
                 finished within that duration.
         """
-        job_id = self._entid(job) if not isinstance(job, str) else job
+        job_id = self._entid(job)
         deadline = (time.monotonic() + timeout) if timeout is not None else None
 
         def _check_timeout() -> None:
             if deadline is not None and time.monotonic() >= deadline:
-                raise TimeoutError(
-                    f"Inference job {job_id} did not finish within {timeout}s"
-                )
+                raise TimeoutError(f"Inference job {job_id} did not finish within {timeout}s")
+
+        def _notify(event: dict) -> None:
+            if on_status is None:
+                return
+            if isinstance(job, InferenceJob):
+                # SSE events are partial updates — apply known fields in-place
+                for key, value in event.items():
+                    try:
+                        setattr(job, key, value)
+                    except Exception:
+                        pass
+                on_status(job)
+            else:
+                on_status(self.get_status(job_id))
 
         # --- Try SSE stream first ---
         try:
             for event in self.stream_status(job_id):
                 _check_timeout()
-                status_str = event.get('status', '')
-                current_job = self._parse_job_response(event)
-                if on_status is not None:
-                    on_status(current_job)
-                if status_str.lower() in _TERMINAL_STATUSES:
-                    return current_job
+                _notify(event)
+                if event.get('status', '').lower() in _TERMINAL_STATUSES:
+                    return
         except Exception as e:
             logger.warning(f"SSE stream ended or failed ({e}); falling back to polling")
 
@@ -213,7 +219,7 @@ def _check_timeout() -> None:
             if on_status is not None:
                 on_status(current_job)
             if current_job.status.lower() in _TERMINAL_STATUSES:
-                return current_job
+                return
             time.sleep(poll_interval)
 
     def cancel(self, job: str | InferenceJob) -> bool:
@@ -392,4 +398,6 @@ def predict_volume(
         )
         response = self._make_request('POST', f'/{self.endpoint_base}/predict-volume', json=payload)
         data = response.json()
-        return self.get_status(data['job_id'])
+        return self.get_status(data['job_id'])
+
+    predict = submit  # Alias for generic prediction endpoint