Merge pull request #764 from superannotateai/FRIDAY_3554

Added ability to filter by item category

Author: VaghinakDev

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

@@ -3421,7 +3421,7 @@ def list_items(
         exclude = {"meta", "annotator_email", "qa_email"}
         if not include_custom_metadata:
             exclude.add("custom_metadata")
-        return BaseSerializer.serialize_iterable(res, exclude=exclude)
+        return BaseSerializer.serialize_iterable(res, exclude=exclude, by_alias=False)
 
     def list_projects(
         self,

src/superannotate/lib/core/entities/__init__.py

@@ -4,6 +4,7 @@
 from lib.core.entities.classes import AnnotationClassEntity
 from lib.core.entities.folder import FolderEntity
 from lib.core.entities.integrations import IntegrationEntity
+from lib.core.entities.items import CategoryEntity
 from lib.core.entities.items import ClassificationEntity
 from lib.core.entities.items import DocumentEntity
 from lib.core.entities.items import ImageEntity
@@ -12,7 +13,6 @@
 from lib.core.entities.items import TiledEntity
 from lib.core.entities.items import VideoEntity
 from lib.core.entities.project import AttachmentEntity
-from lib.core.entities.project import CategoryEntity
 from lib.core.entities.project import ContributorEntity
 from lib.core.entities.project import CustomFieldEntity
 from lib.core.entities.project import ProjectEntity

src/superannotate/lib/core/entities/filters.py

@@ -29,6 +29,8 @@ class ItemFilters(BaseFilters):
     assignments__user_role__in: Optional[List[str]]
     assignments__user_role__ne: Optional[str]
     assignments__user_role__notin: Optional[List[str]]
+    categories__value: Optional[str]
+    categories__value__in: Optional[List[str]]
 
 
 class ProjectFilters(BaseFilters):

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

@@ -2,7 +2,7 @@
 from typing import Optional
 
 from lib.core.entities.base import BaseItemEntity
-from lib.core.entities.base import TimedBaseModel
+from lib.core.entities.project import TimedBaseModel
 from lib.core.enums import ApprovalStatus
 from lib.core.enums import ProjectType
 from lib.core.pydantic_v1 import Extra
@@ -18,9 +18,17 @@ class Config:
         extra = Extra.ignore
 
 
+class CategoryEntity(TimedBaseModel):
+    id: int
+    value: str = Field(None, alias="name")
+
+    class Config:
+        extra = Extra.ignore
+
+
 class MultiModalItemCategoryEntity(TimedBaseModel):
     id: int = Field(None, alias="category_id")
-    name: str = Field(None, alias="category_name")
+    value: str = Field(None, alias="category_name")
 
     class Config:
         extra = Extra.ignore

src/superannotate/lib/core/entities/project.py

@@ -187,8 +187,3 @@ def is_system(self):
 
     class Config:
         extra = Extra.ignore
-
-
-class CategoryEntity(BaseModel):
-    id: Optional[int]
-    name: Optional[str]

src/superannotate/lib/core/serviceproviders.py

@@ -698,6 +698,12 @@ class BaseServiceProvider:
     def get_role_id(self, project: entities.ProjectEntity, role_name: str) -> int:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_category_id(
+        self, project: entities.ProjectEntity, category_name: str
+    ) -> int:
+        raise NotImplementedError
+
     @abstractmethod
     def get_role_name(self, project: entities.ProjectEntity, role_id: int) -> str:
         raise NotImplementedError

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

@@ -2101,8 +2101,10 @@ def execute(self):
                 if categorization_enabled:
                     item_id_category_map = {}
                     for item_name in uploaded_annotations:
-                        category = name_annotation_map[item_name]["metadata"].get(
-                            "item_category"
+                        category = (
+                            name_annotation_map[item_name]["metadata"]
+                            .get("item_category", {})
+                            .get("value")
                         )
                         if category:
                             item_id_category_map[name_item_map[item_name].id] = category

src/superannotate/lib/infrastructure/query_builder.py

@@ -113,6 +113,8 @@ def handle(self, filters: Dict[str, Any], query: Query = None) -> Query:
         for key, val in filters.items():
             _keys = key.split("__")
             val = self._handle_special_fields(_keys, val)
+            if _keys[0] == "categories" and _keys[1] == "value":
+                _keys[1] = "category_id"
             condition, _key = determine_condition_and_key(_keys)
             query &= Filter(_key, val, condition)
         return super().handle(filters, query)
@@ -147,6 +149,14 @@ def _handle_special_fields(self, keys: List[str], val):
                 ]
             else:
                 val = self._service_provider.get_role_id(self._project, val)
+        elif keys[0] == "categories" and keys[1] == "value":
+            if isinstance(val, list):
+                val = [
+                    self._service_provider.get_category_id(self._project, i)
+                    for i in val
+                ]
+            else:
+                val = self._service_provider.get_category_id(self._project, val)
         return val
 
 

src/superannotate/lib/infrastructure/serviceprovider.py

@@ -79,6 +79,13 @@ def list_custom_field_names(self, entity: CustomFieldEntityEnum) -> List[str]:
             self.client.team_id, entity=entity
         )
 
+    def get_category_id(
+        self, project: entities.ProjectEntity, category_name: str
+    ) -> int:
+        return self._cached_work_management_repository.get_category_id(
+            project, category_name
+        )
+
     def get_custom_field_id(
         self, field_name: str, entity: CustomFieldEntityEnum
     ) -> int:

src/superannotate/lib/infrastructure/utils.py

@@ -138,6 +138,23 @@ def get(self, key, **kwargs):
         return self._K_V_map[key]
 
 
+class CategoryCache(BaseCachedWorkManagementRepository):
+    def sync(self, project: ProjectEntity):
+        response = self.work_management.list_project_categories(project.id)
+        if not response.ok:
+            raise AppException(response.error)
+        categories = response.data
+        self._K_V_map[project.id] = {
+            "category_name_id_map": {
+                category.value: category.id for category in categories
+            },
+            "category_id_name_map": {
+                category.id: category.value for category in categories
+            },
+        }
+        self._update_cache_timestamp(project.id)
+
+
 class RoleCache(BaseCachedWorkManagementRepository):
     def sync(self, project: ProjectEntity):
         response = self.work_management.list_workflow_roles(
@@ -221,6 +238,7 @@ def get(self, key, **kwargs):
 
 class CachedWorkManagementRepository:
     def __init__(self, ttl_seconds: int, work_management):
+        self._category_cache = CategoryCache(ttl_seconds, work_management)
         self._role_cache = RoleCache(ttl_seconds, work_management)
         self._status_cache = StatusCache(ttl_seconds, work_management)
         self._project_custom_field_cache = CustomFieldCache(
@@ -236,6 +254,12 @@ def __init__(self, ttl_seconds: int, work_management):
             CustomFieldEntityEnum.TEAM,
         )
 
+    def get_category_id(self, project, category_name: str) -> int:
+        data = self._category_cache.get(project.id, project=project)
+        if category_name in data["category_name_id_map"]:
+            return data["category_name_id_map"][category_name]
+        raise AppException("Invalid category provided.")
+
     def get_role_id(self, project, role_name: str) -> int:
         role_data = self._role_cache.get(project.id, project=project)
         if role_name in role_data["role_name_id_map"]: