updated det/download annoations

Author: VaghinakDev

docs/source/api_reference/api_folder.rst

@@ -8,4 +8,4 @@ Folders
 .. automethod:: superannotate.SAClient.get_folder_by_id
 .. automethod:: superannotate.SAClient.get_folder_metadata
 .. automethod:: superannotate.SAClient.create_folder
-.. automethod:: superannotate.SAClient.delete_folders
+.. automethod:: superannotate.SAClient.delete_folders

src/superannotate/lib/core/serviceproviders.py

@@ -340,6 +340,14 @@ def sort_items_by_size(
     ) -> Dict[str, List]:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_upload_chunks(
+        self,
+        project: entities.ProjectEntity,
+        item_ids: List[int],
+    ) -> Dict[str, List]:
+        raise NotImplementedError
+
     @abstractmethod
     async def download_big_annotation(
         self,

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

@@ -1391,7 +1391,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:
@@ -1440,29 +1439,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:
@@ -1485,26 +1473,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):
@@ -1527,7 +1505,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 = []
@@ -1537,18 +1514,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:
                 nest_asyncio.apply()
                 annotations = asyncio.run(self.run_workers(large_items, small_items))
@@ -1585,7 +1564,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:
@@ -1664,28 +1642,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,
     ):
@@ -1702,19 +1676,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():
@@ -1755,26 +1727,23 @@ 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(
                             large_items, small_items, folder, new_export_path
                         )
                     )
                 except Exception as e:
+                    raise e
                     logger.error(e)
                     self._response.errors = AppException("Can't get annotations.")
                     return self._response

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"
@@ -182,6 +182,26 @@ def sort_items_by_size(
             response_data["large"].extend(response.data.get("large", []))
         return response_data
 
+    def get_upload_chunks(
+        self,
+        project: entities.ProjectEntity,
+        item_ids: List[int],
+    ) -> Dict[str, List]:
+        response_data = {"small": [], "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(
         self,
         project: entities.ProjectEntity,

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/integration/folders/test_set_folder_status.py

@@ -34,9 +34,7 @@ def tearDownClass(cls) -> None:
 
     def test_set_folder_status(self):
         with self.assertLogs("sa", level="INFO") as cm:
-            for index, status in enumerate(
-                self.FOLDER_STATUSES
-            ):
+            for index, status in enumerate(self.FOLDER_STATUSES):
                 sa.set_folder_status(
                     project=self.PROJECT_NAME, folder=self.FOLDER_NAME, status=status
                 )

tests/integration/projects/test_set_project_status.py

@@ -30,9 +30,7 @@ def tearDownClass(cls) -> None:
 
     def test_set_project_status(self):
         with self.assertLogs("sa", level="INFO") as cm:
-            for index, status in enumerate(
-                self.PROJECT_STATUSES
-            ):
+            for index, status in enumerate(self.PROJECT_STATUSES):
                 sa.set_project_status(project=self.PROJECT_NAME, status=status)
                 project = sa.get_project_metadata(self.PROJECT_NAME)
                 assert (

tests/unit/test_async_functions.py

@@ -1,7 +1,7 @@
 import asyncio
+from unittest import TestCase
 
 from superannotate import SAClient
-from unittest import TestCase
 
 
 sa = SAClient()
@@ -30,6 +30,7 @@ def test_get_annotations_in_running_event_loop(self):
         async def _test():
             annotations = sa.get_annotations(self.PROJECT_NAME)
             assert len(annotations) == 4
+
         asyncio.run(_test())
 
     def test_multiple_get_annotations_in_running_event_loop(self):
@@ -38,26 +39,32 @@ async def nested():
             sa.attach_items(self.PROJECT_NAME, self.ATTACH_PAYLOAD)
             annotations = sa.get_annotations(self.PROJECT_NAME)
             assert len(annotations) == 4
+
         async def create_task_test():
             import nest_asyncio
+
             nest_asyncio.apply()
             task1 = asyncio.create_task(nested())
             task2 = asyncio.create_task(nested())
             await task1
             await task2
+
         asyncio.run(create_task_test())
 
         async def gather_test():
             import nest_asyncio
+
             nest_asyncio.apply()
             await asyncio.gather(nested(), nested())
+
         asyncio.run(gather_test())
 
     def test_upload_annotations_in_running_event_loop(self):
         async def _test():
             sa.attach_items(self.PROJECT_NAME, self.ATTACH_PAYLOAD)
-            annotations = sa.upload_annotations(self.PROJECT_NAME, annotations=self.UPLOAD_PAYLOAD)
-            assert len(annotations['succeeded']) == 4
-        asyncio.run(_test())
-
+            annotations = sa.upload_annotations(
+                self.PROJECT_NAME, annotations=self.UPLOAD_PAYLOAD
+            )
+            assert len(annotations["succeeded"]) == 4
 
+        asyncio.run(_test())