Merge pull request #675 from superannotateai/friday_2750

Updated get_annotations to handle project/item ids

Author: VaghinakDev

src/superannotate/lib/app/interface/sdk_interface.py

@@ -2205,20 +2205,28 @@ def invite_contributors_to_team(
         return response.data
 
     def get_annotations(
-        self, project: NotEmptyStr, items: Optional[List[NotEmptyStr]] = None
+        self,
+        project: Union[NotEmptyStr, int],
+        items: Optional[Union[List[NotEmptyStr], List[int]]] = None,
     ):
         """Returns annotations for the given list of items.
 
-        :param project: project name or folder path (e.g., “project1/folder1”).
-        :type project: str
+        :param project: project id or project name or folder path (e.g., “project1/folder1”).
+        :type project: str or int
 
         :param items:  item names. If None, all the items in the specified directory will be used.
-        :type items: list of strs
+        :type items: list of strs or list of ints
 
         :return: list of annotations
-        :rtype: list of strs
+        :rtype: list of dict
         """
-        project, folder = self.controller.get_project_folder_by_path(project)
+        if isinstance(project, str):
+            project, folder = self.controller.get_project_folder_by_path(project)
+        else:
+            project, folder = (
+                self.controller.get_project_by_id(project_id=project).data,
+                None,
+            )
         response = self.controller.annotations.list(project, folder, items)
         if response.errors:
             raise AppException(response.errors)

src/superannotate/lib/core/serviceproviders.py

@@ -246,6 +246,14 @@ def list_by_names(
     ) -> ServiceResponse:
         raise NotImplementedError
 
+    @abstractmethod
+    def list_by_ids(
+        self,
+        project: entities.ProjectEntity,
+        ids: List[int],
+    ) -> ServiceResponse:
+        raise NotImplementedError
+
     @abstractmethod
     def attach(
         self,

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

@@ -22,6 +22,7 @@
 from typing import Optional
 from typing import Set
 from typing import Tuple
+from typing import Union
 
 import aiofiles
 import boto3
@@ -191,10 +192,13 @@ async def upload(_chunk: List[ItemToUpload]):
             [],
         )
         try:
+            items_name_data_map = {i.item.name: i.annotation_json for i in chunk}
+            if not items_name_data_map:
+                return
             response = await service_provider.annotations.upload_small_annotations(
                 project=project,
                 folder=folder,
-                items_name_data_map={i.item.name: i.annotation_json for i in chunk},
+                items_name_data_map=items_name_data_map,
             )
             if response.ok:
                 if response.data.failed_items:  # noqa
@@ -1450,44 +1454,56 @@ def __init__(
         config: ConfigEntity,
         reporter: Reporter,
         project: ProjectEntity,
-        folder: FolderEntity,
-        item_names: Optional[List[str]],
         service_provider: BaseServiceProvider,
+        folder: FolderEntity = None,
+        items: Optional[Union[List[str], List[int]]] = None,
     ):
         super().__init__(reporter)
         self._config = config
         self._project = project
         self._folder = folder
         self._service_provider = service_provider
-        self._item_names = item_names
+        self._items = items
+        self._item_name_id_map = {}
         self._item_names_provided = True
         self._big_annotations_queue = None
 
     def validate_project_type(self):
         if self._project.type == constants.ProjectType.PIXEL.value:
             raise AppException("The function is not supported for Pixel projects.")
 
-    def validate_item_names(self):
-        if self._item_names:  # if names provided
-            unique_item_names = list(set(self._item_names))
-            len_unique_items, len_items = len(unique_item_names), len(self._item_names)
-            if len_unique_items < len_items:
-                self.reporter.log_info(
-                    f"Dropping duplicates. Found {len_unique_items}/{len_items} unique items."
-                )
+    @staticmethod
+    def items_duplication_validation(
+        reporter, items: Union[List[str], List[int]]
+    ) -> Union[List[str], List[int]]:
+        unique_items = list(set(items))
+        len_unique_items, len_items = len(unique_items), len(items)
+        if len_unique_items < len_items:
+            reporter.log_info(
+                f"Dropping duplicates. Found {len_unique_items}/{len_items} unique items."
+            )
+        return unique_items
+
+    def validate_items(self):
+        if self._items:  # if items provided:
+            self.items_duplication_validation(self.reporter, self._items)
             #  Keep order required
             seen = set()
-            self._item_names = [
-                i for i in self._item_names if not (i in seen or seen.add(i))
-            ]
+            self._items = [i for i in self._items if not (i in seen or seen.add(i))]
 
     def _prettify_annotations(self, annotations: List[dict]):
         re_struct = {}
-        if self._item_names:
+        if self._items:
+            names_provided = isinstance(self._items[0], str)
             for annotation in annotations:
-                re_struct[annotation["metadata"]["name"]] = annotation
+                if names_provided:
+                    re_struct[annotation["metadata"]["name"]] = annotation
+                else:
+                    re_struct[
+                        self._item_name_id_map[annotation["metadata"]["name"]]
+                    ] = annotation
             try:
-                return [re_struct[x] for x in self._item_names if x in re_struct]
+                return [re_struct[x] for x in self._items if x in re_struct]
             except KeyError:
                 raise AppException("Broken data.")
 
@@ -1558,18 +1574,28 @@ async def run_workers(
 
     def execute(self):
         if self.is_valid():
-            if self._item_names:
-                items = get_or_raise(
-                    self._service_provider.items.list_by_names(
-                        self._project, self._folder, self._item_names
+            if self._items:
+                if isinstance(self._items[0], str):
+                    items: List[BaseItemEntity] = get_or_raise(
+                        self._service_provider.items.list_by_names(
+                            self._project, self._folder, self._items
+                        )
                     )
-                )
-                len_items, len_provided_items = len(items), len(self._item_names)
+                else:
+                    response = self._service_provider.items.list_by_ids(
+                        project=self._project,
+                        ids=self._items,
+                    )
+                    if not response.ok:
+                        raise AppException(response.error)
+                    items: List[BaseItemEntity] = response.data
+                    self._item_name_id_map = {i.name: i.id for i in items}
+                len_items, len_provided_items = len(items), len(self._items)
                 if len_items != len_provided_items:
                     self.reporter.log_warning(
                         f"Could not find annotations for {len_provided_items - len_items}/{len_provided_items} items."
                     )
-            elif self._item_names is None:
+            elif self._items is None:
                 condition = Condition("project_id", self._project.id, EQ) & Condition(
                     "folder_id", self._folder.id, EQ
                 )
@@ -1583,7 +1609,7 @@ def execute(self):
             items_count = len(items)
             self.reporter.log_info(
                 f"Getting {items_count} annotations from "
-                f"{self._project.name}{f'/{self._folder.name}' if self._folder.name != 'root' else ''}."
+                f"{self._project.name}{f'/{self._folder.name}' if self._folder and self._folder.name != 'root' else ''}."
             )
             id_item_map = {i.id: i for i in items}
             self.reporter.start_progress(
@@ -1602,11 +1628,13 @@ def execute(self):
             try:
                 annotations = run_async(self.run_workers(large_items, small_items))
             except Exception as e:
+                # todo remove
+                raise e
                 logger.error(e)
                 self._response.errors = AppException("Can't get annotations.")
                 return self._response
             self.reporter.finish_progress()
-            self._response.data = self._prettify_annotations(annotations)
+            self._response.data = self._prettify_annotations(annotations)  # noqa
         return self._response
 
 
@@ -1634,15 +1662,10 @@ def __init__(
         self._callback = callback
         self._big_file_queue = None
 
-    def validate_item_names(self):
-        if self._item_names:
-            item_names = list(dict.fromkeys(self._item_names))
-            len_unique_items, len_items = len(item_names), len(self._item_names)
-            if len_unique_items < len_items:
-                self.reporter.log_info(
-                    f"Dropping duplicates. Found {len_unique_items}/{len_items} unique items."
-                )
-                self._item_names = item_names
+    def validate_items(self):
+        self._item_names = GetAnnotations.items_duplication_validation(
+            self.reporter, self._item_names
+        )
 
     @property
     def destination(self) -> str:

src/superannotate/lib/infrastructure/controller.py

@@ -496,16 +496,16 @@ def __init__(self, service_provider: ServiceProvider, config: ConfigEntity):
     def list(
         self,
         project: ProjectEntity,
-        folder: FolderEntity,
-        item_names: List[str],
+        folder: FolderEntity = None,
+        items: Union[List[str], List[int]] = None,
         verbose=True,
     ):
         use_case = usecases.GetAnnotations(
             config=self._config,
             reporter=Reporter(log_info=verbose, log_warning=verbose),
             project=project,
             folder=folder,
-            item_names=item_names,
+            items=items,
             service_provider=self.service_provider,
         )
         return use_case.execute()
@@ -879,8 +879,7 @@ def get_folder_by_id(self, folder_id: int, project_id: int):
     def get_project_by_id(self, project_id: int):
         response = self.projects.get_by_id(project_id=project_id)
         if response.errors:
-            raise AppException(response.errors)
-
+            raise AppException("Project not found.")
         return response
 
     def get_item_by_id(self, item_id: int, project: ProjectEntity):

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

@@ -142,7 +142,7 @@ async def list_small_annotations(
         query_params = {
             "team_id": project.team_id,
             "project_id": project.id,
-            "folder_id": folder.id,
+            # "folder_id": folder.id,
         }
 
         handler = StreamedAnnotations(

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

@@ -4,6 +4,7 @@
 
 from lib.core import entities
 from lib.core.conditions import Condition
+from lib.core.entities import BaseItemEntity
 from lib.core.enums import ProjectType
 from lib.core.exceptions import AppException
 from lib.core.exceptions import BackendError
@@ -12,6 +13,7 @@
 from lib.core.service_types import ImageResponse
 from lib.core.service_types import ItemListResponse
 from lib.core.service_types import PointCloudResponse
+from lib.core.service_types import ServiceResponse
 from lib.core.service_types import TiledResponse
 from lib.core.service_types import VideoResponse
 from lib.core.serviceproviders import BaseItemService
@@ -27,6 +29,7 @@ class ItemService(BaseItemService):
     URL_MOVE_MULTIPLE = "image/move"
     URL_SET_ANNOTATION_STATUSES = "image/updateAnnotationStatusBulk"
     URL_LIST_BY_NAMES = "images/getBulk"
+    URL_LIST_BY_IDS = "images/getImagesByIds"
     URL_COPY_MULTIPLE = "images/copy-image-or-folders"
     URL_COPY_PROGRESS = "images/copy-image-progress"
     URL_DELETE_ITEMS = "image/delete/images"
@@ -75,6 +78,30 @@ def update(self, project: entities.ProjectEntity, item: entities.BaseItemEntity)
             params={"project_id": project.id},
         )
 
+    def list_by_ids(
+        self,
+        project: entities.ProjectEntity,
+        ids: List[int],
+    ):
+        chunk_size = 2000
+        items = []
+        response = None
+        for i in range(0, len(ids), chunk_size):
+            response = self.client.request(
+                self.URL_LIST_BY_IDS,
+                "post",
+                data={
+                    "image_ids": ids[i : i + chunk_size],  # noqa
+                },
+                params={"project_id": project.id},
+                content_type=ServiceResponse,
+            )
+            if not response.ok:
+                return response
+            items.extend(response.data["images"])
+        response.res_data = [BaseItemEntity(**i) for i in items]
+        return response
+
     def list_by_names(
         self,
         project: entities.ProjectEntity,

src/superannotate/lib/infrastructure/stream_data_handler.py

@@ -45,7 +45,9 @@ async def fetch(
         data: dict = None,
         params: dict = None,
     ):
-        kwargs = {"params": params, "json": {"folder_id": params.pop("folder_id")}}
+        kwargs = {"params": params, "json": {}}
+        if "folder_id" in kwargs["params"]:
+            kwargs["json"] = {"folder_id": kwargs["params"].pop("folder_id")}
         if data:
             kwargs["json"].update(data)
         response = await session.request(method, url, **kwargs, timeout=TIMEOUT)  # noqa

tests/integration/annotations/test_get_annotations.py

@@ -42,6 +42,32 @@ def test_get_annotations(self):
                 len(annotation_data["instances"]), len(annotations[0]["instances"])
             )
 
+    def test_get_annotations_by_ids(self):
+        self._attach_items(count=4)  # noqa
+
+        sa.create_annotation_classes_from_classes_json(
+            self.PROJECT_NAME, f"{self.folder_path}/classes/classes.json"
+        )
+        _, _, _ = sa.upload_annotations_from_folder_to_project(
+            self.PROJECT_NAME, self.folder_path
+        )
+        items = sa.search_items(self.PROJECT_NAME)
+
+        annotations = sa.get_annotations(self._project["id"], [i["id"] for i in items])
+
+        self.assertEqual(len(annotations), 4)
+
+    def test_get_annotations_by_wrong_item_ids(self):
+        annotations = sa.get_annotations(self._project["id"], [1, 2, 3])
+
+        self.assertEqual(len(annotations), 0)
+
+    def test_get_annotations_by_wrong_project_ids(self):
+        try:
+            sa.get_annotations(1, [1, 2, 3])
+        except Exception as e:
+            self.assertEqual(str(e), "Project not found.")
+
     @pytest.mark.flaky(reruns=3)
     def test_get_annotations_order(self):
         sa.upload_images_from_folder_to_project(