Merge pull request #799 from superannotateai/FRIDAY-3992

added item category set/remove functions

Author: nareksa

docs/source/api_reference/api_item.rst

@@ -21,3 +21,5 @@ Items
 .. automethod:: superannotate.SAClient.unassign_items
 .. automethod:: superannotate.SAClient.get_item_metadata
 .. automethod:: superannotate.SAClient.set_approval_statuses
+.. automethod:: superannotate.SAClient.set_items_category
+.. automethod:: superannotate.SAClient.remove_items_category

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

@@ -4393,6 +4393,73 @@ def move_items(
             raise AppException(response.errors)
         return response.data
 
+    def set_items_category(
+        self,
+        project: Union[NotEmptyStr, Tuple[int, int], Tuple[str, str]],
+        items: List[Union[int, str]],
+        category: str,
+    ):
+        """
+        Add categories to one or more items.
+
+        :param project: Project and folder as a tuple, folder is optional.
+        :type project: Union[str, Tuple[int, int], Tuple[str, str]]
+
+        :param items: A list of names or IDs of the items to modify.
+        :type items: List[Union[int, str]]
+
+        :param category: Category to assign to the item.
+        :type category: Str
+
+        Request Example:
+        ::
+
+            client.set_items_category(
+                project=("product-review-mm", "folder1"),
+                items=[112233, 112344],
+                category="Shoes"
+            )
+        """
+        project, folder = self.controller.get_project_folder(project)
+        self.controller.check_multimodal_project_categorization(project)
+
+        self.controller.items.attach_detach_items_category(
+            project=project,
+            folder=folder,
+            items=items,
+            category=category,
+            operation="attach",
+        )
+
+    def remove_items_category(
+        self,
+        project: Union[NotEmptyStr, Tuple[int, int], Tuple[str, str]],
+        items: List[Union[int, str]],
+    ):
+        """
+        Remove categories from one or more items.
+
+        :param project: Project and folder as a tuple, folder is optional.
+        :type project: Union[str, Tuple[int, int], Tuple[str, str]]
+
+        :param items: A list of names or IDs of the items to modify.
+        :type items: List[Union[int, str]]
+
+        Request Example:
+        ::
+
+            client.remove_items_category(
+                project=("product-review-mm", "folder1"),
+                items=[112233, 112344]
+            )
+        """
+        project, folder = self.controller.get_project_folder(project)
+        self.controller.check_multimodal_project_categorization(project)
+
+        self.controller.items.attach_detach_items_category(
+            project=project, folder=folder, items=items, operation="detach"
+        )
+
     def set_annotation_statuses(
         self,
         project: Union[NotEmptyStr, dict],

src/superannotate/lib/core/serviceproviders.py

@@ -499,8 +499,14 @@ def delete_multiple(
 
     @abstractmethod
     def bulk_attach_categories(
-        self, project_id: int, folder_id: int, item_category_map: Dict[int, int]
-    ) -> bool:
+        self, project_id: int, folder_id: int, item_id_category_id_map: Dict[int, int]
+    ) -> ServiceResponse:
+        raise NotImplementedError
+
+    @abstractmethod
+    def bulk_detach_categories(
+        self, project_id: int, folder_id: int, item_ids: List[int]
+    ) -> ServiceResponse:
         raise NotImplementedError
 
 

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

@@ -7,6 +7,7 @@
 from typing import Dict
 from typing import Generator
 from typing import List
+from typing import Optional
 from typing import Union
 
 import lib.core as constants
@@ -20,6 +21,7 @@
 from lib.core.entities import ProjectEntity
 from lib.core.entities import VideoEntity
 from lib.core.entities.items import MultiModalItemEntity
+from lib.core.entities.items import ProjectCategoryEntity
 from lib.core.exceptions import AppException
 from lib.core.exceptions import AppValidationException
 from lib.core.exceptions import BackendError
@@ -38,6 +40,7 @@
 from lib.infrastructure.utils import extract_project_folder
 from typing_extensions import Literal
 
+
 logger = logging.getLogger("sa")
 
 
@@ -1272,3 +1275,51 @@ def execute(
             # returning control to the interface function that called it. So no need for
             # error handling in the response
             return self._response
+
+
+class AttacheDetachItemsCategoryUseCase(BaseUseCase):
+    CHUNK_SIZE = 2000
+
+    def __init__(
+        self,
+        project: ProjectEntity,
+        folder: FolderEntity,
+        items: List[MultiModalItemEntity],
+        service_provider: BaseServiceProvider,
+        operation: Literal["attach", "detach"],
+        category: Optional[ProjectCategoryEntity] = None,
+    ):
+        super().__init__()
+        self._project = project
+        self._folder = folder
+        self._items = items
+        self._category = category
+        self._operation = operation
+        self._service_provider = service_provider
+
+    def execute(self):
+        if self._operation == "attach":
+            success_count = 0
+            for chunk in divide_to_chunks(self._items, self.CHUNK_SIZE):
+                item_id_category_id_map: Dict[int, int] = {
+                    i.id: self._category.id for i in chunk
+                }
+                response = self._service_provider.items.bulk_attach_categories(
+                    project_id=self._project.id,
+                    folder_id=self._folder.id,
+                    item_id_category_id_map=item_id_category_id_map,
+                )
+                success_count += len(response.data)
+            logger.info(
+                f"{self._category.name} category successfully added to {success_count} items."
+            )
+        elif self._operation == "detach":
+            success_count = 0
+            for chunk in divide_to_chunks(self._items, self.CHUNK_SIZE):
+                response = self._service_provider.items.bulk_detach_categories(
+                    project_id=self._project.id,
+                    folder_id=self._folder.id,
+                    item_ids=[i.id for i in chunk],
+                )
+                success_count += len(response.data)
+            logger.info(f"Category successfully removed from {success_count} items.")

src/superannotate/lib/infrastructure/controller.py

@@ -35,6 +35,7 @@
 from lib.core.entities.filters import ProjectFilters
 from lib.core.entities.filters import UserFilters
 from lib.core.entities.integrations import IntegrationEntity
+from lib.core.entities.items import ProjectCategoryEntity
 from lib.core.entities.work_managament import ScoreEntity
 from lib.core.entities.work_managament import ScorePayloadEntity
 from lib.core.enums import CustomFieldEntityEnum
@@ -877,7 +878,7 @@ def list_items(
         folder: FolderEntity,
         /,
         include: List[str] = None,
-        **filters: Unpack[ItemFilters],
+        **filters: Optional[Unpack[ItemFilters]],
     ) -> List[BaseItemEntity]:
 
         entity = PROJECT_ITEM_ENTITY_MAP.get(project.type, BaseItemEntity)
@@ -1062,6 +1063,46 @@ def update(self, project: ProjectEntity, item: BaseItemEntity):
         )
         return use_case.execute()
 
+    def attach_detach_items_category(
+        self,
+        project: ProjectEntity,
+        folder: FolderEntity,
+        items: List[Union[int, str]],
+        operation: Literal["attach", "detach"],
+        category: Optional[str] = None,
+    ):
+        if items and isinstance(items[0], str):
+            items = self.list_items(project, folder, name__in=items)
+        elif items and isinstance(items[0], int):
+            items = self.list_items(project, folder, id__in=items)
+        else:
+            raise AppException(
+                "Items must be a list of strings or integers representing item IDs."
+            )
+
+        if category:
+            all_categories = (
+                self.service_provider.work_management.list_project_categories(
+                    project.id, ProjectCategoryEntity  # noqa
+                )
+            )
+            category = next(
+                (c for c in all_categories.data if c.name.lower() == category.lower()),
+                None,
+            )
+            if not category:
+                raise AppException("Category not defined in project.")
+
+        use_case = usecases.AttacheDetachItemsCategoryUseCase(
+            project=project,
+            folder=folder,
+            items=items,
+            category=category,
+            operation=operation,
+            service_provider=self.service_provider,
+        )
+        return use_case.execute()
+
 
 class AnnotationManager(BaseManager):
     def __init__(self, service_provider: ServiceProvider, config: ConfigEntity):

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

@@ -217,7 +217,7 @@ def delete_multiple(self, project: entities.ProjectEntity, item_ids: List[int]):
         )
 
     def bulk_attach_categories(
-        self, project_id: int, folder_id: int, item_category_map: Dict[int, int]
+        self, project_id: int, folder_id: int, item_id_category_id_map: Dict[int, int]
     ) -> bool:
         params = {"project_id": project_id, "folder_id": folder_id}
         response = self.client.request(
@@ -226,10 +226,25 @@ def bulk_attach_categories(
             params=params,
             data={
                 "bulk": [
-                    {"item_id": item_id, "categories": [category]}
-                    for item_id, category in item_category_map.items()
+                    {"item_id": item_id, "categories": [category_id]}
+                    for item_id, category_id in item_id_category_id_map.items()
                 ]
             },
         )
         response.raise_for_status()
-        return response.ok
+        return response
+
+    def bulk_detach_categories(
+        self, project_id: int, folder_id: int, item_ids: List[int]
+    ) -> bool:
+        params = {"project_id": project_id, "folder_id": folder_id}
+        response = self.client.request(
+            self.URL_ATTACH_CATEGORIES,
+            "post",
+            params=params,
+            data={
+                "bulk": [{"item_id": item_id, "categories": []} for item_id in item_ids]
+            },
+        )
+        response.raise_for_status()
+        return response