Merge pull request #588 from superannotateai/1855_annotations

updated get/download annotations

Author: VaghinakDev

src/superannotate/lib/core/serviceproviders.py

@@ -332,10 +332,9 @@ async def list_small_annotations(
         raise NotImplementedError
 
     @abstractmethod
-    def sort_items_by_size(
+    def get_upload_chunks(
         self,
         project: entities.ProjectEntity,
-        folder: entities.FolderEntity,
         item_ids: List[int],
     ) -> Dict[str, List]:
         raise NotImplementedError

src/superannotate/lib/core/usecases/annotations.py

@@ -1387,7 +1387,6 @@ def __init__(
         self._item_names = item_names
         self._item_names_provided = True
         self._big_annotations_queue = None
-        self._small_annotations_queue = None
 
     def validate_project_type(self):
         if self._project.type == constants.ProjectType.PIXEL.value:
@@ -1436,29 +1435,18 @@ async def get_big_annotation(self):
                 break
         return large_annotations
 
-    async def get_small_annotations(self):
-        small_annotations = []
-        while True:
-            items = await self._small_annotations_queue.get()
-            if items:
-                annotations = (
-                    await self._service_provider.annotations.list_small_annotations(
-                        project=self._project,
-                        folder=self._folder,
-                        item_ids=[i.id for i in items],
-                        reporter=self.reporter,
-                    )
-                )
-                small_annotations.extend(annotations)
-            else:
-                await self._small_annotations_queue.put(None)
-                break
-        return small_annotations
+    async def get_small_annotations(self, item_ids: List[int]):
+        return await self._service_provider.annotations.list_small_annotations(
+            project=self._project,
+            folder=self._folder,
+            item_ids=item_ids,
+            reporter=self.reporter,
+        )
 
     async def run_workers(
         self,
         big_annotations: List[BaseItemEntity],
-        small_annotations: List[BaseItemEntity],
+        small_annotations: List[List[Dict]],
     ):
         annotations = []
         if big_annotations:
@@ -1481,26 +1469,16 @@ async def run_workers(
                 )
             )
         if small_annotations:
-            self._small_annotations_queue = asyncio.Queue()
-            small_chunks = divide_to_chunks(
-                small_annotations, size=self._config.ANNOTATION_CHUNK_SIZE
-            )
-            for chunk in small_chunks:
-                self._small_annotations_queue.put_nowait(chunk)
-            self._small_annotations_queue.put_nowait(None)
-
-            annotations.extend(
-                list(
-                    itertools.chain.from_iterable(
-                        await asyncio.gather(
-                            *[
-                                self.get_small_annotations()
-                                for _ in range(self._config.MAX_COROUTINE_COUNT)
-                            ]
-                        )
-                    )
+            for chunks in divide_to_chunks(
+                small_annotations, self._config.MAX_COROUTINE_COUNT
+            ):
+                tasks = []
+                for chunk in chunks:
+                    tasks.append(self.get_small_annotations([i["id"] for i in chunk]))
+                annotations.extend(
+                    list(itertools.chain.from_iterable(await asyncio.gather(*tasks)))
                 )
-            )
+
         return list(filter(None, annotations))
 
     def execute(self):
@@ -1523,7 +1501,6 @@ def execute(self):
                 items = get_or_raise(self._service_provider.items.list(condition))
             else:
                 items = []
-            id_item_map = {i.id: i for i in items}
             if not items:
                 logger.info("No annotations to download.")
                 self._response.data = []
@@ -1533,18 +1510,20 @@ def execute(self):
                 f"Getting {items_count} annotations from "
                 f"{self._project.name}{f'/{self._folder.name}' if self._folder.name != 'root' else ''}."
             )
+            id_item_map = {i.id: i for i in items}
             self.reporter.start_progress(
                 items_count,
                 disable=logger.level > logging.INFO or self.reporter.log_enabled,
             )
-
-            sort_response = self._service_provider.annotations.sort_items_by_size(
-                project=self._project, folder=self._folder, item_ids=list(id_item_map)
+            sort_response = self._service_provider.annotations.get_upload_chunks(
+                project=self._project,
+                item_ids=list(id_item_map),
             )
             large_item_ids = set(map(itemgetter("id"), sort_response["large"]))
-            small_items_ids = set(map(itemgetter("id"), sort_response["small"]))
-            large_items = list(filter(lambda item: item.id in large_item_ids, items))
-            small_items = list(filter(lambda item: item.id in small_items_ids, items))
+            large_items: List[BaseItemEntity] = list(
+                filter(lambda item: item.id in large_item_ids, items)
+            )
+            small_items: List[List[dict]] = sort_response["small"]
             try:
                 annotations = asyncio.run(self.run_workers(large_items, small_items))
             except Exception as e:
@@ -1580,7 +1559,6 @@ def __init__(
         self._service_provider = service_provider
         self._callback = callback
         self._big_file_queue = None
-        self._small_file_queue = None
 
     def validate_item_names(self):
         if self._item_names:
@@ -1659,28 +1637,24 @@ async def download_big_annotations(self, export_path):
                 self._big_file_queue.put_nowait(None)
                 break
 
-    async def download_small_annotations(self, export_path, folder: FolderEntity):
+    async def download_small_annotations(
+        self, item_ids: List[int], export_path, folder: FolderEntity
+    ):
         postfix = self.get_postfix()
-        while True:
-            items = await self._small_file_queue.get()
-            if items:
-                await self._service_provider.annotations.download_small_annotations(
-                    project=self._project,
-                    folder=folder,
-                    item_ids=[i.id for i in items],
-                    reporter=self.reporter,
-                    download_path=f"{export_path}{'/' + self._folder.name if not self._folder.is_root else ''}",
-                    postfix=postfix,
-                    callback=self._callback,
-                )
-            else:
-                self._small_file_queue.put_nowait(None)
-                break
+        await self._service_provider.annotations.download_small_annotations(
+            project=self._project,
+            folder=folder,
+            item_ids=item_ids,
+            reporter=self.reporter,
+            download_path=f"{export_path}{'/' + self._folder.name if not self._folder.is_root else ''}",
+            postfix=postfix,
+            callback=self._callback,
+        )
 
     async def run_workers(
         self,
         big_annotations: List[BaseItemEntity],
-        small_annotations: List[BaseItemEntity],
+        small_annotations: List[List[dict]],
         folder: FolderEntity,
         export_path,
     ):
@@ -1697,19 +1671,17 @@ async def run_workers(
             )
 
         if small_annotations:
-            self._small_file_queue = asyncio.Queue()
-            small_chunks = divide_to_chunks(
-                small_annotations, size=self._config.ANNOTATION_CHUNK_SIZE
-            )
-            for chunk in small_chunks:
-                self._small_file_queue.put_nowait(chunk)
-            self._small_file_queue.put_nowait(None)
-            await asyncio.gather(
-                *[
-                    self.download_small_annotations(export_path, folder)
-                    for _ in range(self._config.MAX_COROUTINE_COUNT)
-                ]
-            )
+            for chunks in divide_to_chunks(
+                small_annotations, self._config.MAX_COROUTINE_COUNT
+            ):
+                tasks = []
+                for chunk in chunks:
+                    tasks.append(
+                        self.download_small_annotations(
+                            [i["id"] for i in chunk], export_path, folder
+                        )
+                    )
+                await asyncio.gather(*tasks)
 
     def execute(self):
         if self.is_valid():
@@ -1749,19 +1721,15 @@ def execute(self):
                     new_export_path += f"/{folder.name}"
 
                 id_item_map = {i.id: i for i in items}
-                sort_response = self._service_provider.annotations.sort_items_by_size(
+                sort_response = self._service_provider.annotations.get_upload_chunks(
                     project=self._project,
-                    folder=self._folder,
                     item_ids=list(id_item_map),
                 )
                 large_item_ids = set(map(itemgetter("id"), sort_response["large"]))
-                small_items_ids = set(map(itemgetter("id"), sort_response["small"]))
-                large_items = list(
+                large_items: List[BaseItemEntity] = list(
                     filter(lambda item: item.id in large_item_ids, items)
                 )
-                small_items = list(
-                    filter(lambda item: item.id in small_items_ids, items)
-                )
+                small_items: List[List[dict]] = sort_response["small"]
                 try:
                     asyncio.run(
                         self.run_workers(

src/superannotate/lib/infrastructure/services/annotation.py

@@ -24,7 +24,7 @@
 
 
 class AnnotationService(BaseAnnotationService):
-    ASSETS_PROVIDER_VERSION = "v2"
+    ASSETS_PROVIDER_VERSION = "v2.01"
     DEFAULT_CHUNK_SIZE = 5000
 
     URL_GET_ANNOTATIONS = "items/annotations/download"
@@ -153,33 +153,24 @@ async def list_small_annotations(
             params=query_params,
         )
 
-    def sort_items_by_size(
+    def get_upload_chunks(
         self,
         project: entities.ProjectEntity,
-        folder: entities.FolderEntity,
         item_ids: List[int],
     ) -> Dict[str, List]:
-        chunk_size = 2000
-        query_params = {
-            "project_id": project.id,
-            "folder_id": folder.id,
-        }
-
         response_data = {"small": [], "large": []}
-        for i in range(0, len(item_ids), chunk_size):
-            body = {
-                "item_ids": item_ids[i : i + chunk_size],  # noqa
-            }  # noqa
-            response = self.client.request(
-                url=urljoin(self.assets_provider_url, self.URL_CLASSIFY_ITEM_SIZE),
-                method="POST",
-                params=query_params,
-                data=body,
-            )
-            if not response.ok:
-                raise AppException(response.error)
-            response_data["small"].extend(response.data.get("small", []))
-            response_data["large"].extend(response.data.get("large", []))
+        response = self.client.request(
+            url=urljoin(self.assets_provider_url, self.URL_CLASSIFY_ITEM_SIZE),
+            method="POST",
+            params={"project_id": project.id, "limit": len(item_ids)},
+            data={"item_ids": item_ids},
+        )
+        if not response.ok:
+            raise AppException(response.error)
+        response_data["small"] = [
+            i["data"] for i in response.data.get("small", {}).values()
+        ]
+        response_data["large"] = response.data.get("large", [])
         return response_data
 
     async def download_big_annotation(

tests/integration/annotations/test_get_annotations.py

@@ -119,7 +119,7 @@ def test_get_annotations10000(self):
             [
                 {"name": f"example_image_{i}.jpg", "url": f"url_{i}"}
                 for i in range(count)
-            ],  # noqa
+            ],
         )
         assert len(sa.search_items(self.PROJECT_NAME)) == count
         a = sa.get_annotations(self.PROJECT_NAME)

tests/unit/test_async_functions.py

@@ -3,7 +3,6 @@
 
 from superannotate import SAClient
 
-
 sa = SAClient()