Added category handling

tod

Author: VaghinakDev

src/superannotate/lib/core/__init__.py

@@ -179,6 +179,7 @@ def setup_logging(level=DEFAULT_LOGGING_LEVEL, file_path=LOG_FILE_LOCATION):
     "Tokenization",
     "ImageAutoAssignEnable",
     "TemplateState",
+    "CategorizeItems",
 ]
 
 __alL__ = (

src/superannotate/lib/core/service_types.py

@@ -226,6 +226,10 @@ class ProjectResponse(ServiceResponse):
     res_data: entities.ProjectEntity = None
 
 
+class ListCategoryResponse(ServiceResponse):
+    res_data: List[entities.CategoryEntity] = None
+
+
 class WorkflowResponse(ServiceResponse):
     res_data: entities.WorkflowEntity = None
 

src/superannotate/lib/core/serviceproviders.py

@@ -9,13 +9,13 @@
 
 from lib.core import entities
 from lib.core.conditions import Condition
-from lib.core.entities import CategoryEntity
 from lib.core.jsx_conditions import Query
 from lib.core.reporter import Reporter
 from lib.core.service_types import AnnotationClassListResponse
 from lib.core.service_types import FolderListResponse
 from lib.core.service_types import FolderResponse
 from lib.core.service_types import IntegrationListResponse
+from lib.core.service_types import ListCategoryResponse
 from lib.core.service_types import ProjectListResponse
 from lib.core.service_types import ProjectResponse
 from lib.core.service_types import ServiceResponse
@@ -89,13 +89,13 @@ def list_workflow_roles(self, project_id: int, workflow_id: int):
         raise NotImplementedError
 
     @abstractmethod
-    def list_project_categories(self, project_id: int) -> List[entities.CategoryEntity]:
+    def list_project_categories(self, project_id: int) -> ListCategoryResponse:
         raise NotImplementedError
 
     @abstractmethod
     def create_project_categories(
-        self, project_id: int, categories: List[CategoryEntity]
-    ):
+        self, project_id: int, categories: List[str]
+    ) -> ServiceResponse:
         raise NotImplementedError
 
 
@@ -362,7 +362,9 @@ def delete_multiple(
         raise NotImplementedError
 
     @abstractmethod
-    def bulk_attach_categories(self, project_id: int, folder_id: int, item_category_map: Dict[int, int]) -> bool:
+    def bulk_attach_categories(
+        self, project_id: int, folder_id: int, item_category_map: Dict[int, int]
+    ) -> bool:
         raise NotImplementedError
 
 

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

@@ -10,6 +10,7 @@
 import time
 import traceback
 from collections import defaultdict
+from contextlib import suppress
 from dataclasses import dataclass
 from itertools import islice
 from operator import itemgetter
@@ -1846,6 +1847,7 @@ def __init__(
         self._transform_version = (
             "llmJsonV2" if transform_version is None else transform_version
         )
+        self._category_name_to_id_map = {}
 
     @property
     def files_queue(self):
@@ -1994,11 +1996,12 @@ def serialize_folder_name(val):
 
     def execute(self):
         if self.is_valid():
+            # TODO check categories status in the project
+            skip_categorization = False
             serialized_original_folder_map = {}
             failed, skipped, uploaded = [], [], []
-            distributed_items: Dict[str, Dict[str, Any]] = defaultdict(
-                dict
-            )  # folder_id -> item_name -> annotation
+            # folder_id -> item_name -> annotation
+            distributed_items: Dict[str, Dict[str, Any]] = defaultdict(dict)
             valid_items_count = 0
             for annotation in self._annotations:
                 if self._validate_json(annotation):
@@ -2068,6 +2071,17 @@ def execute(self):
                     {i.item.name for i in items_to_upload}
                     - set(failed_annotations).union(skipped)
                 )
+                if not skip_categorization:
+                    item_id_category_map = {}
+                    for item_name in uploaded_annotations:
+                        category = name_annotation_map[item_name]["metadata"].get(
+                            "item_category"
+                        )
+                        if category:
+                            item_id_category_map[name_item_map[item_name].id] = category
+                    self._attach_categories(
+                        folder_id=folder.id, item_id_category_map=item_id_category_map
+                    )
                 workflow = self._service_provider.work_management.get_workflow(
                     self._project.workflow_id
                 )
@@ -2096,3 +2110,44 @@ def execute(self):
                 "skipped": skipped,
             }
             return self._response
+
+    def _attach_categories(self, folder_id: int, item_id_category_map: Dict[int, str]):
+        categories_to_create: List[str] = []
+        item_id_category_id_map: Dict[int, int] = {}
+        if not self._category_name_to_id_map:
+            response = self._service_provider.work_management.list_project_categories(
+                self._project.id
+            )
+            response.raise_for_status()
+            categories = response.data
+            self._category_name_to_id_map = {c.name: c.id for c in categories}
+        for item_id in list(item_id_category_map.keys()):
+            category_name = item_id_category_map[item_id]
+            if category_name not in self._category_name_to_id_map:
+                categories_to_create.append(category_name)
+            else:
+                item_id_category_id_map[item_id] = self._category_name_to_id_map[
+                    category_name
+                ]
+                item_id_category_map.pop(item_id)
+
+        if categories_to_create:
+            _categories = (
+                self._service_provider.work_management.create_project_categories(
+                    project_id=self._project.id,
+                    categories=categories_to_create,
+                ).data["data"]
+            )
+            for c in _categories:
+                self._category_name_to_id_map[c["name"]] = c["id"]
+            for item_id, category_name in item_id_category_map.items():
+                with suppress(KeyError):
+                    item_id_category_id_map[item_id] = self._category_name_to_id_map[
+                        category_name
+                    ]
+        if item_id_category_id_map:
+            self._service_provider.items.bulk_attach_categories(
+                project_id=self._project.id,
+                folder_id=folder_id,
+                item_category_map=item_id_category_id_map,
+            )

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

@@ -21,7 +21,7 @@ class ItemService(BaseItemService):
     URL_DELETE_ITEMS = "image/delete/images"
     URL_SET_APPROVAL_STATUSES = "/items/bulk/change"
     URL_COPY_MOVE_MULTIPLE = "images/copy-move-images-folders"
-    URL_ATTACH_CATEGORIES #
+    URL_ATTACH_CATEGORIES = "items/bulk/setcategory"
 
     def update(self, project: entities.ProjectEntity, item: entities.BaseItemEntity):
         return self.client.request(
@@ -214,12 +214,20 @@ def delete_multiple(self, project: entities.ProjectEntity, item_ids: List[int]):
             data={"image_ids": item_ids},
         )
 
-
-    def bulk_attach_categories(self, project_id: int, folder_id: int, item_category_map: Dict[int, int]) -> bool:
-        params = {
-            "project_id": project_id,
-            "folder_id": folder_id
-        }
+    def bulk_attach_categories(
+        self, project_id: int, folder_id: int, item_category_map: Dict[int, int]
+    ) -> bool:
+        params = {"project_id": project_id, "folder_id": folder_id}
         response = self.client.request(
-            self.
-        )
+            self.URL_ATTACH_CATEGORIES,
+            "post",
+            params=params,
+            data={
+                "bulk": [
+                    {"item_id": item_id, "categories": [category]}
+                    for item_id, category in item_category_map.items()
+                ]
+            },
+        )
+        response.raise_for_status()
+        return response.ok

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

@@ -8,6 +8,8 @@
 from lib.core.jsx_conditions import Filter
 from lib.core.jsx_conditions import OperatorEnum
 from lib.core.jsx_conditions import Query
+from lib.core.service_types import ListCategoryResponse
+from lib.core.service_types import ServiceResponse
 from lib.core.serviceproviders import BaseWorkManagementService
 
 
@@ -26,7 +28,7 @@ def _generate_context(**kwargs):
         encoded_context = base64.b64encode(json.dumps(kwargs).encode("utf-8"))
         return encoded_context.decode("utf-8")
 
-    def list_project_categories(self, project_id: int) -> List[CategoryEntity]:
+    def list_project_categories(self, project_id: int) -> ListCategoryResponse:
         return self.client.paginate(
             self.URL_LIST_CATEGORIES,
             item_type=CategoryEntity,
@@ -39,15 +41,21 @@ def list_project_categories(self, project_id: int) -> List[CategoryEntity]:
         )
 
     def create_project_categories(
-        self, project_id: int, categories: List[CategoryEntity]
-    ):
+        self, project_id: int, categories: List[str]
+    ) -> ServiceResponse:
         response = self.client.request(
-            "post",
+            method="post",
             url=self.URL_CREATE_CATEGORIES,
-            pararms={"project_id": project_id},
-            data={"bulk": [i["name"] for i in categories]},
+            params={"project_id": project_id},
+            data={"bulk": [{"name": i} for i in categories]},
+            headers={
+                "x-sa-entity-context": self._generate_context(
+                    team_id=self.client.team_id, project_id=project_id
+                ),
+            },
         )
         response.raise_for_status()
+        return response
 
     def get_workflow(self, pk: int) -> WorkflowEntity:
         response = self.list_workflows(Filter("id", pk, OperatorEnum.EQ))

tests/data_set/multimodal/annotations/jsonl/form1_with_categories.jsonl

@@ -0,0 +1,3 @@
+{"metadata":  {"folder_name": "test_folder", "name":  "0123456789101112", "item_category":  "category1"}, ",data": {"component_id_0":  {"value":  ["Partially complete, needs review"]}, "component_id_1": {"value":  "I am a text input 001"}, "component_id_2": {"value": 11}}}
+{"metadata":  {"folder_name":  " test_Folder", "name": "item_002", "item_category":  "category2"}, "data": {"component_id_0":  {"value":  ["Partially complete, needs review", "Incomplete"]}, "component_id_1": {"value":  "I am a text input 002"}, "component_id_2": {"value":  33}}}
+{"metadata":  {"folder_name":  "Test_Folder ", "name": "item_003", "item_category":  "category3"}, "data": {"component_id_0":  {"value": ["Partially complete, needs review", "Incomplete"]}, "component_id_1": {"value":  "I am a text input 003"}, "component_id_2": {"value":  33}}}

tests/integration/annotations/test_upload_annotations.py

@@ -143,6 +143,9 @@ class MultiModalUploadAnnotations(BaseTestCase):
     JSONL_ANNOTATIONS_PATH = os.path.join(
         DATA_SET_PATH, "multimodal/annotations/jsonl/form1.jsonl"
     )
+    JSONL_ANNOTATIONS_WITH_CATEGORIES_PATH = os.path.join(
+        DATA_SET_PATH, "multimodal/annotations/jsonl/form1_with_categories.jsonl"
+    )
     CLASSES_TEMPLATE_PATH = os.path.join(
         Path(__file__).parent.parent.parent,
         "data_set/editor_templates/from1_classes.json",
@@ -154,10 +157,13 @@ def setUp(self, *args, **kwargs):
             self.PROJECT_NAME,
             self.PROJECT_DESCRIPTION,
             "Multimodal",
-            settings=[{"attribute": "TemplateState", "value": 1}],
+            settings=[
+                {"attribute": "CategorizeItems", "value": 1},
+                {"attribute": "TemplateState", "value": 1},
+            ],
         )
         project = sa.controller.get_project(self.PROJECT_NAME)
-        time.sleep(5)
+        time.sleep(2)
         with open(self.EDITOR_TEMPLATE_PATH) as f:
             res = sa.controller.service_provider.projects.attach_editor_template(
                 sa.controller.team, project, template=json.load(f)
@@ -195,3 +201,10 @@ def test_error_upload_from_folder_to_folder_(self):
                 "You can't include a folder when uploading from within a folder.",
             ):
                 sa.upload_annotations(f"{self.PROJECT_NAME}/tmp", annotations=data)
+
+    def test_upload_with_categories(self):
+        with open(self.JSONL_ANNOTATIONS_WITH_CATEGORIES_PATH) as f:
+            data = [json.loads(line) for line in f]
+            sa.upload_annotations(f"{self.PROJECT_NAME}", annotations=data)
+            annotations = sa.get_annotations(f"{self.PROJECT_NAME}/test_folder")
+            assert len(annotations) == 3