Feat/incremental update dat 643

datamint/apihandler/annotation_api_handler.py

@@ -13,6 +13,9 @@
 from .dto.annotation_dto import CreateAnnotationDto, LineGeometry, BoxGeometry, CoordinateSystem, AnnotationType
 import pydicom
 import json
+from deprecated import deprecated
+from pathlib import Path
+from tqdm.auto import tqdm
 
 _LOGGER = logging.getLogger(__name__)
 _USER_LOGGER = logging.getLogger('user_logger')
@@ -267,8 +270,9 @@ async def _upload_volume_segmentation_async(self,
             raise NotImplementedError("`name=string` is not supported yet for volume segmentation.")
         if isinstance(name, dict):
             if any(isinstance(k, tuple) for k in name.keys()):
-                raise NotImplementedError("For volume segmentations, `name` must be a dictionary with integer keys only.") 
-
+                raise NotImplementedError(
+                    "For volume segmentations, `name` must be a dictionary with integer keys only.")
+
         # Prepare file for upload
         if isinstance(file_path, str):
             if file_path.endswith('.nii') or file_path.endswith('.nii.gz'):
@@ -1098,6 +1102,29 @@ def delete_annotation(self, annotation_id: str | dict):
         resp = self._run_request(request_params)
         self._check_errors_response_json(resp)
 
+    def get_annotation_by_id(self, annotation_id: str) -> dict:
+        """
+        Get an annotation by its unique id.
+
+        Args:
+            annotation_id (str): The annotation unique id.
+
+        Returns:
+            dict: The annotation information.
+        """
+        request_params = {
+            'method': 'GET',
+            'url': f'{self.root_url}/annotations/{annotation_id}',
+        }
+
+        try:
+            resp = self._run_request(request_params)
+            return resp.json()
+        except HTTPError as e:
+            _LOGGER.error(f"Error getting annotation by id {annotation_id}: {e}")
+            raise
+
+    @deprecated(reason="Use download_segmentation_file instead")
     def get_segmentation_file(self, resource_id: str, annotation_id: str) -> bytes:
         request_params = {
             'method': 'GET',
@@ -1107,6 +1134,35 @@ def get_segmentation_file(self, resource_id: str, annotation_id: str) -> bytes:
         resp = self._run_request(request_params)
         return resp.content
 
+    def download_segmentation_file(self, annotation: str | dict, fpath_out: str | Path | None) -> bytes:
+        """
+        Download the segmentation file for a given resource and annotation.
+
+        Args:
+            annotation (str | dict): The annotation unique id or an annotation object.
+            fpath_out (str | None): (Optional) The file path to save the downloaded segmentation file.
+
+        Returns:
+            bytes: The content of the downloaded segmentation file in bytes format.
+        """
+        if isinstance(annotation, dict):
+            annotation_id = annotation['id']
+            resource_id = annotation['resource_id']
+        else:
+            annotation_id = annotation
+            resource_id = self.get_annotation_by_id(annotation_id)['resource_id']
+
+        request_params = {
+            'method': 'GET',
+            'url': f'{self.root_url}/annotations/{resource_id}/annotations/{annotation_id}/file',
+        }
+
+        resp = self._run_request(request_params)
+        if fpath_out is not None:
+            with open(str(fpath_out), 'wb') as f:
+                f.write(resp.content)
+        return resp.content
+
     def set_annotation_status(self,
                               project_id: str,
                               resource_id: str,
@@ -1124,3 +1180,69 @@ def set_annotation_status(self,
         }
         resp = self._run_request(request_params)
         self._check_errors_response_json(resp)
+
+
+    async def _async_download_segmentation_file(self,
+                                                annotation: str | dict,
+                                                save_path: str | Path,
+                                                session: aiohttp.ClientSession | None = None,
+                                                progress_bar: tqdm | None = None):
+        """
+        Asynchronously download a segmentation file.
+
+        Args:
+            annotation (str | dict): The annotation unique id or an annotation object.
+            save_path (str | Path): The path to save the file.
+            session (aiohttp.ClientSession): The aiohttp session to use for the request.
+            progress_bar (tqdm | None): Optional progress bar to update after download completion.
+        """
+        if isinstance(annotation, dict):
+            annotation_id = annotation['id']
+            resource_id = annotation['resource_id']
+        else:
+            annotation_id = annotation
+            # TODO: This is inefficient as it requires an extra API call per annotation
+            # Consider passing resource_id separately or caching annotation info
+            resource_id = self.get_annotation_by_id(annotation_id)['resource_id']
+
+        url = f'{self.root_url}/annotations/{resource_id}/annotations/{annotation_id}/file'
+        request_params = {
+            'method': 'GET',
+            'url': url
+        }
+
+        try:
+            data_bytes = await self._run_request_async(request_params, session, 'content')
+            with open(save_path, 'wb') as f:
+                f.write(data_bytes)
+            if progress_bar:
+                progress_bar.update(1)
+        except ResourceNotFoundError as e:
+            e.set_params('annotation', {'annotation_id': annotation_id})
+            raise e
+
+    def download_multiple_segmentations(self,
+                                        annotations: list[str | dict],
+                                        save_paths: list[str | Path] | str
+                                        ) -> None:
+        """
+        Download multiple segmentation files and save them to the specified paths.
+
+        Args:
+            annotations (list[str | dict]): A list of annotation unique ids or annotation objects.
+            save_paths (list[str | Path] | str): A list of paths to save the files or a directory path.
+        """
+        async def _download_all_async():
+            async with aiohttp.ClientSession() as session:
+                tasks = [
+                    self._async_download_segmentation_file(annotation, save_path=path, session=session, progress_bar=progress_bar)
+                    for annotation, path in zip(annotations, save_paths)
+                ]
+                await asyncio.gather(*tasks)
+
+        if isinstance(save_paths, str):
+            save_paths = [os.path.join(save_paths, f"{ann['id'] if isinstance(ann, dict) else ann}") for ann in annotations]
+
+        with tqdm(total=len(annotations), desc="Downloading segmentations", unit="file") as progress_bar:
+            loop = asyncio.get_event_loop()
+            loop.run_until_complete(_download_all_async())

datamint/apihandler/base_api_handler.py

@@ -85,7 +85,7 @@ def __init__(self,
             msg = f"API key not provided! Use the environment variable " + \
                 f"{BaseAPIHandler.DATAMINT_API_VENV_NAME} or pass it as an argument."
             raise DatamintException(msg)
-        self.semaphore = asyncio.Semaphore(10)  # Limit to 10 parallel requests
+        self.semaphore = asyncio.Semaphore(20)
 
         if check_connection:
             self.check_connection()
@@ -157,30 +157,34 @@ def _generate_curl_command(self, request_args: dict) -> str:
     async def _run_request_async(self,
                                  request_args: dict,
                                  session: aiohttp.ClientSession | None = None,
-                                 data_to_get: str = 'json'):
+                                 data_to_get: Literal['json', 'text', 'content'] = 'json'):
         if session is None:
             async with aiohttp.ClientSession() as s:
-                return await self._run_request_async(request_args, s)
-        try:
-            _LOGGER.debug(f"Running request to {request_args['url']}")
-            _LOGGER.debug(f'Equivalent curl command: "{self._generate_curl_command(request_args)}"')
-        except Exception as e:
-            _LOGGER.debug(f"Error generating curl command: {e}")
-
-        # add apikey to the headers
-        if 'headers' not in request_args:
-            request_args['headers'] = {}
-
-        request_args['headers']['apikey'] = self.api_key
-
-        async with session.request(**request_args) as response:
-            self._check_errors_response(response, request_args)
-            if data_to_get == 'json':
-                return await response.json()
-            elif data_to_get == 'text':
-                return await response.text()
-            else:
-                raise ValueError("data_to_get must be either 'json' or 'text'")
+                return await self._run_request_async(request_args, s, data_to_get)
+
+        async with self.semaphore:
+            try:
+                _LOGGER.debug(f"Running request to {request_args['url']}")
+                _LOGGER.debug(f'Equivalent curl command: "{self._generate_curl_command(request_args)}"')
+            except Exception as e:
+                _LOGGER.debug(f"Error generating curl command: {e}")
+
+            # add apikey to the headers
+            if 'headers' not in request_args:
+                request_args['headers'] = {}
+
+            request_args['headers']['apikey'] = self.api_key
+
+            async with session.request(**request_args) as response:
+                self._check_errors_response(response, request_args)
+                if data_to_get == 'json':
+                    return await response.json()
+                elif data_to_get == 'text':
+                    return await response.text()
+                elif data_to_get == 'content':
+                    return await response.read()
+                else:
+                    raise ValueError("data_to_get must be either 'json' or 'text'")
 
     def _check_errors_response(self,
                                response,
@@ -237,9 +241,9 @@ def _get_endpoint_url(self, endpoint: str) -> str:
         return f'{self.root_url}/{endpoint}'
 
     def _run_pagination_request(self,
-                                request_params: Dict,
-                                return_field: Optional[Union[str, List]] = None
-                                ) -> Generator[Dict, None, None]:
+                                request_params: dict,
+                                return_field: str | list | None = None
+                                ) -> Generator[dict | list, None, None]:
         offset = 0
         params = request_params.get('params', {})
         while True: