Merge pull request #495 from superannotateai/get_annotations

get annotations uses assets provider

Author: VaghinakDev

src/superannotate/lib/core/serviceproviders.py

@@ -421,30 +421,32 @@ def list_annotation_classes(
 
     @abstractmethod
     async def download_big_annotation(
-            self,
-            project_id: int,
-            team_id: int,
-            download_path: str,
-            postfix: str,
-            item: dict,
-            callback: Callable = None,
+        self,
+        project_id: int,
+        team_id: int,
+        download_path: str,
+        postfix: str,
+        item: dict,
+        callback: Callable = None,
     ):
         raise NotImplementedError
 
     @abstractmethod
     async def download_small_annotations(
-            self,
-            project_id: int,
-            team_id: int,
-            folder_id: int,
-            reporter: Reporter,
-            download_path: str,
-            postfix: str,
-            items: List[str] = None,
-            callback: Callable = None,
+        self,
+        project_id: int,
+        team_id: int,
+        folder_id: int,
+        reporter: Reporter,
+        download_path: str,
+        postfix: str,
+        items: List[str] = None,
+        callback: Callable = None,
     ):
         raise NotImplementedError
 
     @abstractmethod
-    def sort_items_by_size(self, item_names: List[str], team_id: int, project_id: int, folder_id: int):
-        raise NotImplementedError
+    def sort_items_by_size(
+        self, item_names: List[str], team_id: int, project_id: int, folder_id: int
+    ):
+        raise NotImplementedError

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

@@ -775,36 +775,95 @@ def validate_item_names(self):
             self._item_names = [item.name for item in self._images.get_all(condition)]
 
     def _prettify_annotations(self, annotations: List[dict]):
+        restruct = {}
+
         if self._item_names_provided:
+            for annotation in annotations:
+                restruct[annotation["metadata"]["name"]] = annotation
             try:
-                data = []
-                for annotation in annotations:
-                    data.append(
-                        (
-                            self._item_names.index(annotation["metadata"]["name"]),
-                            annotation,
-                        )
-                    )
-                return [i[1] for i in sorted(data, key=lambda x: x[0])]
+                return [restruct[x] for x in self._item_names if x in restruct]
             except KeyError:
                 raise AppException("Broken data.")
+
+        return annotations
+
+    async def get_big_annotation(
+        self,
+    ):
+
+        large_annotations = []
+        while True:
+            item = await self._big_annotations_queue.get()
+            if not item:
+                await self._big_annotations_queue.put(None)
+                break
+
+            large_annotation = await self._client.get_big_annotation(
+                team_id=self._project.team_id,
+                project_id=self._project.id,
+                folder_id=self._folder.uuid,
+                item=item,
+                reporter=self.reporter,
+            )
+
+            large_annotations.append(large_annotation)
+
+        return large_annotations
+
+    async def get_small_annotations(self, item_names):
+        small_annotations = await self._client.get_small_annotations(
+            team_id=self._project.team_id,
+            project_id=self._project.id,
+            folder_id=self._folder.uuid,
+            items=item_names,
+            reporter=self.reporter,
+        )
+
+        return small_annotations
+
+    async def distribute_to_queue(self, big_annotations):
+
+        for item in big_annotations:
+            await self._big_annotations_queue.put(item)
+
+        await self._big_annotations_queue.put(None)
+
+    async def run_workers(self, big_annotations, small_annotations):
+
+        annotations = await asyncio.gather(
+            self.distribute_to_queue(big_annotations),
+            self.get_small_annotations(small_annotations),
+            self.get_big_annotation(),
+            self.get_big_annotation(),
+            self.get_big_annotation(),
+            return_exceptions=True,
+        )
+
+        annotations = [i for x in annotations[1:] for i in x if x]
         return annotations
 
     def execute(self):
         if self.is_valid():
+            self._big_annotations_queue = asyncio.Queue()
             items_count = len(self._item_names)
             self.reporter.log_info(
                 f"Getting {items_count} annotations from "
                 f"{self._project.name}{f'/{self._folder.name}' if self._folder.name != 'root' else ''}."
             )
             self.reporter.start_progress(items_count, disable=not self._show_process)
-            annotations = self._client.get_annotations(
+
+            items = self._client.sort_items_by_size(
+                item_names=self._item_names,
                 team_id=self._project.team_id,
-                project_id=self._project.id,
                 folder_id=self._folder.uuid,
-                items=self._item_names,
-                reporter=self.reporter,
+                project_id=self._project.id,
+            )
+
+            small_annotations = [x["name"] for x in items["small"]]
+            annotations = asyncio.run(
+                self.run_workers(items["large"], small_annotations)
             )
+
             received_items_count = len(annotations)
             self.reporter.finish_progress()
             if items_count > received_items_count:
@@ -1202,7 +1261,6 @@ def execute(self):
                     if not folder.is_root and self._folder.is_root:
                         new_export_path += f"/{folder.name}"
 
-
                     # TODO check
                     if not item_names:
                         continue