From c7a8ec9a10df2fd4515975503197b1e1717d41e3 Mon Sep 17 00:00:00 2001 From: Narek Mkhitaryan Date: Fri, 4 Jul 2025 15:12:05 +0400 Subject: [PATCH] added item category set/remove functions --- docs/source/api_reference/api_item.rst | 2 + .../lib/app/interface/sdk_interface.py | 67 +++++++++ .../lib/core/serviceproviders.py | 10 +- src/superannotate/lib/core/usecases/items.py | 51 +++++++ .../lib/infrastructure/controller.py | 43 +++++- .../lib/infrastructure/services/item.py | 23 ++- .../integration/items/test_attach_category.py | 142 ++++++++++++++++++ 7 files changed, 331 insertions(+), 7 deletions(-) create mode 100644 tests/integration/items/test_attach_category.py diff --git a/docs/source/api_reference/api_item.rst b/docs/source/api_reference/api_item.rst index 0947ed7b..e25f5789 100644 --- a/docs/source/api_reference/api_item.rst +++ b/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 diff --git a/src/superannotate/lib/app/interface/sdk_interface.py b/src/superannotate/lib/app/interface/sdk_interface.py index 856a3f2a..dec26352 100644 --- a/src/superannotate/lib/app/interface/sdk_interface.py +++ b/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], diff --git a/src/superannotate/lib/core/serviceproviders.py b/src/superannotate/lib/core/serviceproviders.py index 95b313f1..6672383b 100644 --- a/src/superannotate/lib/core/serviceproviders.py +++ b/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 diff --git a/src/superannotate/lib/core/usecases/items.py b/src/superannotate/lib/core/usecases/items.py index 10fc9e00..fe9a6491 100644 --- a/src/superannotate/lib/core/usecases/items.py +++ b/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.") diff --git a/src/superannotate/lib/infrastructure/controller.py b/src/superannotate/lib/infrastructure/controller.py index ae361ee3..dd531a32 100644 --- a/src/superannotate/lib/infrastructure/controller.py +++ b/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): diff --git a/src/superannotate/lib/infrastructure/services/item.py b/src/superannotate/lib/infrastructure/services/item.py index d99d37b5..828802cb 100644 --- a/src/superannotate/lib/infrastructure/services/item.py +++ b/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 diff --git a/tests/integration/items/test_attach_category.py b/tests/integration/items/test_attach_category.py new file mode 100644 index 00000000..eb812db3 --- /dev/null +++ b/tests/integration/items/test_attach_category.py @@ -0,0 +1,142 @@ +import json +import os +import time +from pathlib import Path +from unittest import TestCase + +from src.superannotate import SAClient + +sa = SAClient() + + +class TestItemAttachCategory(TestCase): + PROJECT_NAME = "TestItemAttachCategory" + PROJECT_TYPE = "Multimodal" + PROJECT_DESCRIPTION = "DESCRIPTION" + EDITOR_TEMPLATE_PATH = os.path.join( + Path(__file__).parent.parent.parent, + "data_set/editor_templates/form1.json", + ) + CLASSES_TEMPLATE_PATH = os.path.join( + Path(__file__).parent.parent.parent, + "data_set/editor_templates/form1_classes.json", + ) + + @classmethod + def setUpClass(cls, *args, **kwargs) -> None: + cls.tearDownClass() + cls._project = sa.create_project( + cls.PROJECT_NAME, + cls.PROJECT_DESCRIPTION, + cls.PROJECT_TYPE, + settings=[ + {"attribute": "TemplateState", "value": 1}, + {"attribute": "CategorizeItems", "value": 1}, + ], + ) + team = sa.controller.team + project = sa.controller.get_project(cls.PROJECT_NAME) + time.sleep(5) + + with open(cls.EDITOR_TEMPLATE_PATH) as f: + template_data = json.load(f) + res = sa.controller.service_provider.projects.attach_editor_template( + team, project, template=template_data + ) + assert res.ok + sa.create_annotation_classes_from_classes_json( + cls.PROJECT_NAME, cls.CLASSES_TEMPLATE_PATH + ) + + @classmethod + def tearDownClass(cls) -> None: + # cleanup test scores and project + projects = sa.search_projects(cls.PROJECT_NAME, return_metadata=True) + for project in projects: + try: + sa.delete_project(project) + except Exception: + pass + + @staticmethod + def _attach_items(path: str, names: list[str]): + sa.attach_items(path, [{"name": name, "url": f"url-{name}"} for name in names]) + + def test_attache_category(self): + self._attach_items(self.PROJECT_NAME, ["item-1", "item-2"]) + sa.create_categories(self.PROJECT_NAME, ["category-1", "category-2"]) + + with self.assertLogs("sa", level="INFO") as cm: + sa.set_items_category(self.PROJECT_NAME, ["item-1", "item-2"], "category-1") + assert ( + "INFO:sa:category-1 category successfully added to 2 items." + == cm.output[0] + ) + + items = sa.list_items(self.PROJECT_NAME, include=["categories"]) + assert all(i["categories"][0]["value"] == "category-1" for i in items) + + def test_remove_items_category(self): + self._attach_items(self.PROJECT_NAME, ["item-1", "item-2", "item-3"]) + sa.create_categories(self.PROJECT_NAME, ["category-1", "category-2"]) + sa.set_items_category( + self.PROJECT_NAME, ["item-1", "item-2", "item-3"], "category-1" + ) + + items = sa.list_items(self.PROJECT_NAME, include=["categories"]) + assert len(items) == 3 + assert all( + len(i["categories"]) == 1 and i["categories"][0]["value"] == "category-1" + for i in items + ) + + with self.assertLogs("sa", level="INFO") as cm: + sa.remove_items_category(self.PROJECT_NAME, ["item-1", "item-2"]) + assert "INFO:sa:Category successfully removed from 2 items." == cm.output[0] + + items = sa.list_items(self.PROJECT_NAME, include=["categories"]) + item_dict = {item["name"]: item for item in items} + + assert len(item_dict["item-1"]["categories"]) == 0 + assert len(item_dict["item-2"]["categories"]) == 0 + + assert len(item_dict["item-3"]["categories"]) == 1 + assert item_dict["item-3"]["categories"][0]["value"] == "category-1" + + def test_remove_items_category_by_ids(self): + self._attach_items(self.PROJECT_NAME, ["item-4", "item-5"]) + sa.create_categories(self.PROJECT_NAME, ["category-test"]) + sa.set_items_category(self.PROJECT_NAME, ["item-4", "item-5"], "category-test") + + items = sa.list_items(self.PROJECT_NAME, include=["categories"]) + item_ids = [ + item["id"] for item in items if item["name"] in ["item-4", "item-5"] + ] + + sa.remove_items_category(self.PROJECT_NAME, item_ids) + items = sa.list_items(self.PROJECT_NAME, include=["categories"]) + for item in items: + if item["name"] in ["item-4", "item-5"]: + assert len(item["categories"]) == 0 + + def test_remove_items_category_with_folder(self): + folder_name = "test-folder" + sa.create_folder(self.PROJECT_NAME, folder_name) + folder_path = f"{self.PROJECT_NAME}/{folder_name}" + self._attach_items(folder_path, ["folder-item-1", "folder-item-2"]) + + sa.create_categories(self.PROJECT_NAME, ["folder-category"]) + sa.set_items_category( + folder_path, ["folder-item-1", "folder-item-2"], "folder-category" + ) + + sa.remove_items_category(folder_path, ["folder-item-1"]) + + items = sa.list_items( + project=self.PROJECT_NAME, folder=folder_name, include=["categories"] + ) + item_dict = {item["name"]: item for item in items} + + assert len(item_dict["folder-item-1"]["categories"]) == 0 + assert len(item_dict["folder-item-2"]["categories"]) == 1 + assert item_dict["folder-item-2"]["categories"][0]["value"] == "folder-category"