Add keypoint handling

Author: VaghinakDev

CHANGELOG.rst

@@ -7,7 +7,7 @@ History
 All release highlights of this project will be documented in this file.
 
 4.4.34 - April 11, 2025
-______________________
+_______________________
 
 **Added**
 

docs/source/api_reference/api_project.rst

@@ -24,6 +24,6 @@ Projects
 .. automethod:: superannotate.SAClient.add_contributors_to_project
 .. automethod:: superannotate.SAClient.get_project_settings
 .. automethod:: superannotate.SAClient.set_project_default_image_quality_in_editor
-.. automethod:: superannotate.SAClient.set_project_steps
 .. automethod:: superannotate.SAClient.get_project_steps
+.. automethod:: superannotate.SAClient.set_project_steps
 .. automethod:: superannotate.SAClient.get_component_config

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

@@ -74,7 +74,6 @@
 from lib.core.entities.work_managament import WMUserTypeEnum
 from lib.core.jsx_conditions import EmptyQuery
 
-
 logger = logging.getLogger("sa")
 
 NotEmptyStr = constr(strict=True, min_length=1)
@@ -1485,10 +1484,11 @@ def get_project_steps(self, project: Union[str, dict]):
         :param project: project name or metadata
         :type project: str or dict
 
-        :return: project steps
-        :rtype: list of dicts
+        :return: A list of step dictionaries,
+            or a dictionary containing both steps and their connections (for Keypoint workflows).
+        :rtype: list of dicts or dict
 
-        Response Example:
+        Response Example for General Annotation Project:
         ::
 
             [
@@ -1507,6 +1507,34 @@ def get_project_steps(self, project: Union[str, dict]):
                 }
             ]
 
+        Response Example for Keypoint Annotation Project:
+        ::
+
+            {
+              "steps": [
+                {
+                  "step": 1,
+                  "className": "Left Shoulder",
+                  "class_id": "1",
+                  "attribute": [
+                    {
+                            "attribute": {
+                                "id": 123,
+                                "group_id": 12
+                            }
+                        }
+                  ]
+                },
+                {
+                  "step": 2,
+                  "class_id": "2",
+                  "className": "Right Shoulder",
+                }
+              ],
+              "connections": [
+                [1, 2]
+              ]
+            }
         """
         project_name, _ = extract_project_folder(project)
         project = self.controller.get_project(project_name)
@@ -2503,7 +2531,12 @@ def download_export(
         if response.errors:
             raise AppException(response.errors)
 
-    def set_project_steps(self, project: Union[NotEmptyStr, dict], steps: List[dict]):
+    def set_project_steps(
+        self,
+        project: Union[NotEmptyStr, dict],
+        steps: List[dict],
+        connections: List[List[int]] = None,
+    ):
         """Sets project's steps.
 
         :param project: project name or metadata
@@ -2512,7 +2545,11 @@ def set_project_steps(self, project: Union[NotEmptyStr, dict], steps: List[dict]
         :param steps: new workflow list of dicts
         :type steps: list of dicts
 
-        Request Example:
+        :param connections: Defines connections between keypoint annotation steps.
+            Each inner list specifies a pair of step IDs indicating a connection.
+        :type connections: list of dicts
+
+        Request Example for General Annotation Project:
         ::
 
             sa.set_project_steps(
@@ -2533,10 +2570,40 @@ def set_project_steps(self, project: Union[NotEmptyStr, dict], steps: List[dict]
                     }
                 ]
             )
+
+        Request Example for Keypoint Annotation Project:
+        ::
+
+            sa.set_project_steps(
+                project="Pose Estimation Project",
+                steps=[
+                    {
+                        "step": 1,
+                        "class_id": 12,
+                        "attribute": [
+                            {
+                                "attribute": {
+                                    "id": 123,
+                                    "group_id": 12
+                                }
+                            }
+                        ]
+                    },
+                    {
+                        "step": 2,
+                        "class_id": 13
+                    }
+                ],
+                connections=[
+                    [1, 2]
+                ]
+            )
         """
         project_name, _ = extract_project_folder(project)
         project = self.controller.get_project(project_name)
-        response = self.controller.projects.set_steps(project, steps=steps)
+        response = self.controller.projects.set_steps(
+            project, steps=steps, connections=connections
+        )
         if response.errors:
             raise AppException(response.errors)
 

src/superannotate/lib/core/enums.py

@@ -116,6 +116,7 @@ class ProjectType(BaseTitledEnum):
     def images(self):
         return self.VECTOR.value, self.PIXEL.value, self.TILED.value
 
+
 class StepsType(Enum):
     INITIAL = 1
     BASIC = 2

src/superannotate/lib/core/serviceproviders.py

@@ -273,7 +273,7 @@ def set_step(self, project: entities.ProjectEntity, step: entities.StepEntity):
         raise NotImplementedError
 
     @abstractmethod
-    def set_keypoint_steps(self, project: entities.ProjectEntity, steps):
+    def set_keypoint_steps(self, project: entities.ProjectEntity, steps, connections):
         raise NotImplementedError
 
     @abstractmethod

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

@@ -21,7 +21,6 @@
 from lib.core.usecases.base import BaseUseCase
 from lib.core.usecases.base import BaseUserBasedUseCase
 
-
 logger = logging.getLogger("sa")
 
 
@@ -491,8 +490,12 @@ def validate_project_type(self):
 
     def execute(self):
         if self.is_valid():
-            project_settings = self._service_provider.projects.list_settings(project=self._project).data
-            step_setting = next((i for i in project_settings if i.attribute == "WorkflowType"), None)
+            project_settings = self._service_provider.projects.list_settings(
+                project=self._project
+            ).data
+            step_setting = next(
+                (i for i in project_settings if i.attribute == "WorkflowType"), None
+            )
             if step_setting.value == constants.StepsType.BASIC:
                 data = []
                 steps = self._service_provider.projects.list_steps(self._project).data
@@ -508,9 +511,11 @@ def execute(self):
                     data.append(step_data)
                 self._response.data = data
             else:
-                steps = self._service_provider.projects.list_keypoint_steps(self._project).data
-                raise NotImplementedError
-
+                self._response.data = (
+                    self._service_provider.projects.list_keypoint_steps(
+                        self._project
+                    ).data["steps"]
+                )
         return self._response
 
 
@@ -586,17 +591,40 @@ def __init__(
         service_provider: BaseServiceProvider,
         steps: list,
         project: ProjectEntity,
+        connections: List[List[int]] = None,
     ):
         super().__init__()
         self._service_provider = service_provider
         self._steps = steps
+        self._connections = connections
         self._project = project
 
     def validate_project_type(self):
         if self._project.type in constants.LIMITED_FUNCTIONS:
             raise AppValidationException(
                 constants.LIMITED_FUNCTIONS[self._project.type]
             )
+
+    def validate_connections(self):
+        if not self._connections:
+            return
+
+        if len(self._connections) > len(self._steps):
+            raise AppValidationException(
+                "Invalid connections: more connections than steps."
+            )
+
+        possible_connections = set(range(1, len(self._steps) + 1))
+        for connection_group in self._connections:
+            if len(set(connection_group)) != len(connection_group):
+                raise AppValidationException(
+                    "Invalid connections: duplicates in a connection group."
+                )
+            if not set(connection_group).issubset(possible_connections):
+                raise AppValidationException(
+                    "Invalid connections: index out of allowed range."
+                )
+
     def set_basic_steps(self, annotation_classes):
         annotation_classes_map = {}
         annotations_classes_attributes_map = {}
@@ -618,9 +646,7 @@ def set_basic_steps(self, annotation_classes):
             project=self._project,
             steps=self._steps,
         )
-        existing_steps = self._service_provider.projects.list_steps(
-            self._project
-        ).data
+        existing_steps = self._service_provider.projects.list_steps(self._project).data
         existing_steps_map = {}
         for steps in existing_steps:
             existing_steps_map[steps.step] = steps.id
@@ -630,12 +656,10 @@ def set_basic_steps(self, annotation_classes):
             annotation_class_name = step["className"]
             for attribute in step["attribute"]:
                 attribute_name = attribute["attribute"]["name"]
-                attribute_group_name = attribute["attribute"]["attribute_group"][
-                    "name"
-                ]
+                attribute_group_name = attribute["attribute"]["attribute_group"]["name"]
                 if not annotations_classes_attributes_map.get(
-                        f"{annotation_class_name}__{attribute_group_name}__{attribute_name}",
-                        None,
+                    f"{annotation_class_name}__{attribute_group_name}__{attribute_name}",
+                    None,
                 ):
                     raise AppException(
                         f"Attribute group name or attribute name not found {attribute_group_name}."
@@ -656,10 +680,44 @@ def set_basic_steps(self, annotation_classes):
             attributes=req_data,
         )
 
-    def set_keypoint_steps(self, annotation_classes):
+    @staticmethod
+    def _validate_keypoint_steps(annotation_classes, steps):
+        class_group_attrs_map = {}
+        for annotation_class in annotation_classes:
+            class_group_attrs_map[annotation_class.id] = dict()
+            for group in annotation_class.attribute_groups:
+                class_group_attrs_map[annotation_class.id][group.id] = []
+                for attribute in group.attributes:
+                    class_group_attrs_map[annotation_class.id][group.id].append(
+                        attribute.id
+                    )
+        for step in steps:
+            class_id = step.get("class_id", None)
+            if not class_id or class_id not in class_group_attrs_map:
+                raise AppException("Annotation class not found.")
+            attributes = step.get("attribute", None)
+            if not attributes:
+                continue
+            for attr in attributes:
+                try:
+                    _id, group_id = attr["attribute"].get("id", None), attr[
+                        "attribute"
+                    ].get("group_id", None)
+                    assert _id in class_group_attrs_map[class_id][group_id]
+                except (KeyError, AssertionError):
+                    raise AppException("Invalid steps provided.")
+
+    def set_keypoint_steps(self, annotation_classes, steps, connections):
+        self._validate_keypoint_steps(annotation_classes, steps)
+        for i in range(1, len(self._steps) + 1):
+            step = self._steps[i - 1]
+            step["id"] = i
+            if "attribute" not in step:
+                step["attribute"] = []
         self._service_provider.projects.set_keypoint_steps(
             project=self._project,
-            steps=self._steps,
+            steps=steps,
+            connections=connections,
         )
 
     def execute(self):
@@ -669,13 +727,24 @@ def execute(self):
                 Condition("project_id", self._project.id, EQ)
             ).data
 
-            project_settings = self._service_provider.projects.list_settings(project=self._project).data
-            step_setting = next((i for i in project_settings if i.attribute == "WorkflowType"), None)
+            project_settings = self._service_provider.projects.list_settings(
+                project=self._project
+            ).data
+            step_setting = next(
+                (i for i in project_settings if i.attribute == "WorkflowType"), None
+            )
+            if (
+                self._connections is not None
+                and step_setting.value == constants.StepsType.BASIC.value
+            ):
+                raise AppException("Can't update steps type. ")
 
             if step_setting.value == constants.StepsType.BASIC:
                 self.set_basic_steps(annotation_classes)
             else:
-                self.set_keypoint_steps(annotation_classes)
+                self.set_keypoint_steps(
+                    annotation_classes, self._steps, self._connections
+                )
 
         return self._response
 

src/superannotate/lib/infrastructure/controller.py

@@ -489,10 +489,13 @@ def list_steps(self, project: ProjectEntity):
         )
         return use_case.execute()
 
-    def set_steps(self, project: ProjectEntity, steps: List):
+    def set_steps(
+        self, project: ProjectEntity, steps: List, connections: List[List[int]] = None
+    ):
         use_case = usecases.SetStepsUseCase(
             service_provider=self.service_provider,
             steps=steps,
+            connections=connections,
             project=project,
         )
         return use_case.execute()

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

@@ -107,10 +107,7 @@ def list_steps(self, project: entities.ProjectEntity):
         )
 
     def list_keypoint_steps(self, project: entities.ProjectEntity):
-        return self.client.request(
-            self.URL_KEYPOINT_STEPS.format(project.id),
-            "get"
-        )
+        return self.client.request(self.URL_KEYPOINT_STEPS.format(project.id), "get")
 
     def set_step(self, project: entities.ProjectEntity, step: entities.StepEntity):
         return self.client.request(
@@ -119,11 +116,16 @@ def set_step(self, project: entities.ProjectEntity, step: entities.StepEntity):
             data={"steps": [step]},
         )
 
-    def set_keypoint_steps(self, project: entities.ProjectEntity, steps):
+    def set_keypoint_steps(self, project: entities.ProjectEntity, steps, connections):
         return self.client.request(
             self.URL_SET_KEYPOINT_STEPS.format(project.id),
             "post",
-            data={"steps": steps},
+            data={
+                "steps": {
+                    "steps": steps,
+                    "connections": connections if connections else [],
+                }
+            },
         )
 
     # TODO check