Merge pull request #637 from superannotateai/relationship_class_2089

add Relationship class type support

Author: nareksa

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

@@ -55,6 +55,8 @@
 from lib.core.entities.classes import AttributeGroup
 from lib.core.entities.integrations import IntegrationEntity
 from lib.core.enums import ImageQuality
+from lib.core.enums import ProjectType
+from lib.core.enums import ClassTypeEnum
 from lib.core.exceptions import AppException
 from lib.core.types import MLModel
 from lib.core.types import PriorityScoreEntity
@@ -72,11 +74,16 @@
 PROJECT_STATUS = Literal["NotStarted", "InProgress", "Completed", "OnHold"]
 
 PROJECT_TYPE = Literal[
-    "Vector", "Pixel", "Video", "Document", "Tiled", "Other", "PointCloud", "CustomEditor"
+    "Vector",
+    "Pixel",
+    "Video",
+    "Document",
+    "Tiled",
+    "Other",
+    "PointCloud",
+    "CustomEditor",
 ]
 
-CLASS_TYPE = Literal["object", "tag"]
-
 ANNOTATION_STATUS = Literal[
     "NotStarted", "InProgress", "QualityCheck", "Returned", "Completed", "Skipped"
 ]
@@ -1310,7 +1317,7 @@ def create_annotation_class(
         name: NotEmptyStr,
         color: NotEmptyStr,
         attribute_groups: Optional[List[AttributeGroup]] = None,
-        class_type: CLASS_TYPE = "object",
+        class_type: str = "object",
     ):
         """Create annotation class in project
 
@@ -1331,7 +1338,7 @@ def create_annotation_class(
               -  "name"
         :type attribute_groups: list of dicts
 
-        :param class_type: class type. Should be either "object" or "tag"
+        :param class_type: class type. Should be either "object" or "tag". Document project type can also have "relationship" type of classes.
         :type class_type: str
 
         :return: new class metadata
@@ -1405,6 +1412,13 @@ def create_annotation_class(
         except ValidationError as e:
             raise AppException(wrap_error(e))
         project = self.controller.projects.get_by_name(project).data
+        if (
+            project.type != ProjectType.DOCUMENT
+            and annotation_class.type == ClassTypeEnum.RELATIONSHIP
+        ):
+            raise AppException(
+                f"{annotation_class.type.name} class type is not supported in {project.type.name} project."
+            )
         response = self.controller.annotation_classes.create(
             project=project, annotation_class=annotation_class
         )

src/superannotate/lib/core/enums.py

@@ -157,6 +157,7 @@ class AnnotationStatus(BaseTitledEnum):
 class ClassTypeEnum(BaseTitledEnum):
     OBJECT = "object", 1
     TAG = "tag", 2
+    RELATIONSHIP = "relationship", 3
 
     @classmethod
     def get_value(cls, name):

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

@@ -573,7 +573,10 @@ def get_annotation_from_s3(bucket, path: str):
 
     def prepare_annotation(self, annotation: dict, size) -> dict:
         errors = None
-        if size < BIG_FILE_THRESHOLD and self._project.type < constants.ProjectType.PIXEL.value:
+        if (
+            size < BIG_FILE_THRESHOLD
+            and self._project.type < constants.ProjectType.PIXEL.value
+        ):
             use_case = ValidateAnnotationUseCase(
                 reporter=self.reporter,
                 team_id=self._project.team_id,

tests/data_set/document_annotation/classes/classes.json

@@ -28,5 +28,16 @@
         ]
       }
     ]
+  },
+  {
+    "id": 5175192,
+    "project_id": 555040,
+    "name": "apple",
+    "color": "#bafbcd",
+    "count": 0,
+    "type": "relationship",
+    "createdAt": "2023-08-24T14:45:13.000Z",
+    "updatedAt": "2023-08-24T14:45:13.000Z",
+    "attribute_groups": []
   }
 ]

tests/integration/annotations/test_download_annotations.py

@@ -131,3 +131,50 @@ def test_download_annotations_duplicated_names(self):
                 assert (
                     "INFO:sa:Dropping duplicates. Found 1/4 unique items." in cm.output
                 )
+
+
+class TestDocumentProjectDownloadAnnotations(BaseTestCase):
+    PROJECT_NAME = "Test-document-project-download_annotations"
+    FOLDER_NAME = "FOLDER_NAME"
+    FOLDER_NAME_2 = "FOLDER_NAME_2"
+    PROJECT_DESCRIPTION = "Desc"
+    PROJECT_TYPE = "Document"
+    TEST_FOLDER_PATH = "data_set/document_annotation"
+    ITEM_NAME = "text_file_example_1"
+
+    @property
+    def folder_path(self):
+        return os.path.join(Path(__file__).parent.parent.parent, self.TEST_FOLDER_PATH)
+
+    def test_document_project_download_annotations(self):
+        sa.attach_items(
+            self.PROJECT_NAME,
+            [
+                {
+                    "name": self.ITEM_NAME,
+                    "url": "https://sa-public-files.s3.us-west-2.amazonaws.com/Text+project/text_file_example_1.txt",
+                }
+            ],
+        )
+        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
+        )
+        with tempfile.TemporaryDirectory() as temp_dir:
+            annotations_path = sa.download_annotations(
+                f"{self.PROJECT_NAME}", temp_dir, [self.ITEM_NAME]
+            )
+            self.assertEqual(len(os.listdir(temp_dir)), 2)
+            with open(
+                f"{self.folder_path}//{self.ITEM_NAME}.json"
+            ) as pre_annotation_file, open(
+                f"{annotations_path}/{self.ITEM_NAME}.json"
+            ) as post_annotation_file:
+                pre_annotation_data = json.load(pre_annotation_file)
+                post_annotation_data = json.load(post_annotation_file)
+                self.assertEqual(
+                    len(pre_annotation_data["instances"]),
+                    len(post_annotation_data["instances"]),
+                )

tests/integration/classes/test_create_annotation_class.py

@@ -382,3 +382,24 @@ def test_create_annotation_class_with_default_attribute(self):
                     }
                 ],
             )
+
+
+class TestDocumentAnnotationClasses(BaseTestCase):
+    PROJECT_NAME = "TestDocumentAnnotationClasses"
+    PROJECT_DESCRIPTION = "desc"
+    PROJECT_TYPE = "Document"
+
+    def test_document_project_create_annotation_class_search(self):
+        sa.create_annotation_class(
+            self.PROJECT_NAME, name="tt", color="#FFFFFF", class_type="relationship"
+        )
+        classes = sa.search_annotation_classes(self.PROJECT_NAME)
+        self.assertEqual(len(classes), 1)
+        self.assertEqual(classes[0]["type"], "relationship")
+        self.assertEqual(classes[0]["color"], "#FFFFFF")
+        sa.create_annotation_class(self.PROJECT_NAME, name="tb", color="#FFFFFF")
+        #  test search
+        classes = sa.search_annotation_classes(self.PROJECT_NAME, "bb")
+        self.assertEqual(len(classes), 0)
+        classes = sa.search_annotation_classes(self.PROJECT_NAME, "tt")
+        self.assertEqual(len(classes), 1)