ITEP-67565: Allow users to sort labels in any order they want (#239)

Author: JortBergfeld

interactive_ai/libs/iai_core_py/iai_core/entities/graph.py

@@ -129,7 +129,3 @@ class MultiDiGraph(Graph):
 
     def __init__(self) -> None:
         super().__init__(directed=True)
-
-    def topological_sort(self):
-        """Returns a generator of nodes in topologically sorted order."""
-        return nx.topological_sort(self._graph)

interactive_ai/libs/iai_core_py/iai_core/entities/label.py

@@ -145,14 +145,7 @@ def __repr__(self):
     def __eq__(self, other: object) -> bool:
         """Returns True if the two labels are equal."""
         if isinstance(other, Label):
-            return (
-                self.id_ == other.id_
-                and self.name == other.name
-                and self.color == other.color
-                and self.hotkey == other.hotkey
-                and self.domain == other.domain
-                and self.is_anomalous == other.is_anomalous
-            )
+            return self.id_ == other.id_
         return False
 
     def __lt__(self, other: object):

interactive_ai/libs/iai_core_py/iai_core/entities/label_schema.py

@@ -13,7 +13,6 @@
 
 from iai_core.entities.graph import MultiDiGraph
 from iai_core.entities.label import Label
-from iai_core.entities.scored_label import ScoredLabel
 from iai_core.utils.uid_generator import generate_uid
 
 from geti_types import ID, PersistentEntity
@@ -73,15 +72,10 @@ def __init__(
     ):
         self.id_ = ID(ObjectId()) if id is None else id
 
-        self.labels = sorted(labels, key=natural_sort_label_id)
+        self.labels = list(labels)
         self.name = name
         self.group_type = group_type
 
-    @property
-    def minimum_label_id(self) -> ID:
-        """Returns the minimum (oldest) label ID, which is the first label in self.labels since this list is sorted."""
-        return self.labels[0].id_
-
     def remove_label(self, label: Label) -> None:
         """Remove label from label group if it exists in the group.
 
@@ -103,7 +97,7 @@ def __eq__(self, other: object):
         """Returns True if the LabelGroup is equal to the other object."""
         if not isinstance(other, LabelGroup):
             return False
-        return self.id_ == other.id_ and (set(self.labels) == set(other.labels) and self.group_type == other.group_type)
+        return self.id_ == other.id_
 
     def __repr__(self) -> str:
         """Returns the string representation of the LabelGroup."""
@@ -119,8 +113,6 @@ class LabelTree(MultiDiGraph):
     def __init__(self) -> None:
         super().__init__()
 
-        self.__topological_order_cache: list[Label] | None = None
-
     def add_edge(self, node1: Label, node2: Label, edge_value: Any = None) -> None:
         """Add edge between two nodes in the tree.
 
@@ -129,50 +121,24 @@ def add_edge(self, node1: Label, node2: Label, edge_value: Any = None) -> None:
         :param edge_value: The value of the new edge. Defaults to None.
         """
         super().add_edge(node1, node2, edge_value)
-        self.clear_topological_cache()
 
     def add_node(self, node: Label) -> None:
         """Add node to the tree."""
         super().add_node(node)
-        self.clear_topological_cache()
 
     def add_edges(self, edges: Any) -> None:
         """Add edges between Labels."""
         self._graph.add_edges_from(edges)
-        self.clear_topological_cache()
 
     def remove_node(self, node: Label) -> None:
         """Remove node from the tree."""
         super().remove_node(node)
-        self.clear_topological_cache()
 
     @property
     def num_labels(self) -> int:
         """Return the number of labels in the tree."""
         return self.num_nodes()
 
-    def clear_topological_cache(self) -> None:
-        """Clear the internal cache of the list of labels sorted in topological order.
-
-        This function should be called if the topology of the graph has changed to
-            prevent the cache from being stale.
-        Note that it is automatically called when modifying the topology through the
-            methods provided by this class.
-        """
-        self.__topological_order_cache = None
-
-    def get_labels_in_topological_order(self) -> list[Label]:
-        """Return a list of the labels in this graph sorted in topological order.
-
-        To avoid performance issues, the output of this function is cached.
-        """
-        if self.__topological_order_cache is None:
-            # TODO: It seems that we are storing the edges the wrong way around.
-            #       To work around this issue, we have to reverse the sorted list.
-            self.__topological_order_cache = list(reversed(list(self.topological_sort())))
-
-        return self.__topological_order_cache
-
     @property
     def type(self) -> str:
         """Returns the type of the LabelTree."""
@@ -181,7 +147,6 @@ def type(self) -> str:
     def add_child(self, parent: Label, child: Label) -> None:
         """Add a `child` Label to `parent`."""
         self.add_edge(child, parent)
-        self.clear_topological_cache()
 
     def get_parent(self, label: Label) -> Label | None:
         """Returns the parent of `label`"""
@@ -322,12 +287,12 @@ def get_labels(self, include_empty: bool) -> list[Label]:
         :param include_empty: flag determining whether to include empty labels
         :return: list of labels in the label schema
         """
-        labels = []
-        for group in self._groups:
-            for label in group.labels:
-                if (include_empty or not label.is_empty) and label.id_ not in self.deleted_label_ids:
-                    labels.append(label)
-        return sorted(labels, key=lambda x: x.id_)
+        return [
+            label
+            for group in self._groups
+            for label in group.labels
+            if (include_empty or not label.is_empty) and label.id_ not in self.deleted_label_ids
+        ]
 
     def get_label_map(self) -> dict[ID, Label]:
         """
@@ -345,7 +310,7 @@ def get_empty_labels(self) -> tuple[Label, ...]:
 
         :return: tuple of empty labels in the label schema
         """
-        return tuple(sorted([label for label in self.get_labels(include_empty=True) if label.is_empty]))
+        return tuple(label for label in self.get_labels(include_empty=True) if label.is_empty)
 
     def get_label_ids(self, include_empty: bool) -> list[ID]:
         """
@@ -364,9 +329,7 @@ def get_all_labels(self) -> list[Label]:
 
         :return: list of labels in the label schema
         """
-        labels = [label for group in self._groups for label in group.labels]
-
-        return sorted(labels, key=lambda x: x.id_)
+        return [label for group in self._groups for label in group.labels]
 
     def get_groups(self, include_empty: bool = False) -> list[LabelGroup]:
         """
@@ -594,14 +557,7 @@ def __repr__(self) -> str:
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, LabelSchema):
             return False
-        return (
-            self.id_ == other.id_
-            and self.project_id == other.project_id
-            and self.previous_schema_revision_id == other.previous_schema_revision_id
-            and self.label_tree == other.label_tree
-            and self.get_groups(include_empty=True) == other.get_groups(include_empty=True)
-            and self.deleted_label_ids == other.deleted_label_ids
-        )
+        return self.id_ == other.id_
 
 
 class NullLabelSchema(LabelSchema):
@@ -676,7 +632,7 @@ def from_parent(
         label_groups = []
 
         for parent_group in parent_schema.get_groups(include_empty=True):
-            group_labels = list(set(parent_group.labels).intersection(set_of_labels))
+            group_labels = [label for label in parent_group.labels if label in set_of_labels]
             if len(group_labels) > 0:
                 label_groups.append(
                     LabelGroup(
@@ -737,32 +693,4 @@ def __repr__(self) -> str:
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, LabelSchemaView):
             return False
-        return (
-            self.parent_schema == other.parent_schema
-            and self.task_node_id == other.task_node_id
-            and self.previous_schema_revision_id == other.previous_schema_revision_id
-            and self.label_tree == other.label_tree
-            and self.get_groups(include_empty=True) == other.get_groups(include_empty=True)
-        )
-
-
-def natural_sort_label_id(target: ID | Label | ScoredLabel) -> list[int | str]:
-    """Generates a natural sort key for a Label object based on its ID.
-
-    Example:
-    origin_sorted_labels = sorted(labels, key=lambda x: x.id_)
-    natural_sorted_labels = sorted(labels, key=lambda x: x.natural_sort_label_id)
-
-    print(origin_sorted_labels) # Output: [Label(0), Label(1), Label(10), ... Label(2)]
-    print(natural_sorted_labels) # Output: [Label(0), Label(1), Label(2), ... Label(10)]
-
-    :param target (Union[ID, Label]): The ID or Label or ScoredLabel object to be sorted.
-    :returns: List[Union[int, str]]: A list of integers representing the numeric substrings in the ID
-    in the order they appear.
-    """
-
-    if isinstance(target, Label | ScoredLabel):
-        target = target.id_
-    if isinstance(target, str) and target.isdecimal():
-        return ["", int(target)]  # "" is added for the case where id of some lables is None
-    return [target]
+        return self.id_ == other.id_

interactive_ai/libs/iai_core_py/iai_core/repos/mappers/mongodb_mappers/label_mapper.py

@@ -91,13 +91,13 @@ def backward(instance: dict, project_identifier: ProjectIdentifier) -> LabelGrou
         from iai_core.repos import LabelRepo
 
         label_repo = LabelRepo(project_identifier)
+        label_ids = [IDToMongo.backward(label_id) for label_id in instance["label_ids"]]
+        label_map = label_repo.get_by_ids(label_ids)
         return LabelGroup(
             id=IDToMongo.backward(instance["_id"]),
             name=instance["name"],
             group_type=LabelGroupType[instance["relation_type"]],
-            labels=list(
-                label_repo.get_by_ids([IDToMongo.backward(label_id) for label_id in instance["label_ids"]]).values()
-            ),
+            labels=[label_map[label_id] for label_id in label_ids],
         )
 
 

interactive_ai/libs/iai_core_py/iai_core/utils/project_builder.py

@@ -985,38 +985,65 @@ def edit_existing_project(  # noqa: C901, PLR0912, PLR0915
                 if label_parent:
                     child_to_parent_id[label] = label_parent
                 group_name_by_label[label.id_] = label_group
-            labels_by_task[task_node.id_] = updated_labels
+
+            # Sort labels according to ordering in request
+            ordered_label_names = parser.get_custom_labels_names_by_task(task_name=task_name)
+            ordered_updated_labels = []
+            for label_name in ordered_label_names:
+                for label in updated_labels:
+                    if label_name == label.name:
+                        ordered_updated_labels.append(label)
+
+            # Add empty label, if it exists in the old labels, but not in the updated ones
+            if not any(label.is_empty for label in ordered_updated_labels):
+                empty_label = next((label for label in old_labels if label.is_empty), None)
+                if empty_label is not None:
+                    ordered_updated_labels.append(empty_label)
+                    empty_group = next(group for group in old_groups if group.group_type == LabelGroupType.EMPTY_LABEL)
+                    group_name_by_label[empty_label.id_] = empty_group.name
+
+            labels_by_task[task_node.id_] = ordered_updated_labels
 
             if is_keypoint_detection_enabled and task_node.task_properties.task_type == TaskType.KEYPOINT_DETECTION:
                 keypoint_structure_data = parser.get_keypoint_structure_data(task_name=task_name)
                 keypoint_structure = cls._build_keypoint_structure(
                     keypoint_structure_data=keypoint_structure_data,
-                    labels=updated_labels,
+                    labels=ordered_updated_labels,
                 )
                 project.keypoint_structure = keypoint_structure
 
-        if labels_structure_changed:
-            # Labels structure changed, so the label_schema and task_to_label_schema_view are recomputed from scratch.
-            (
-                label_schema,
-                task_to_label_schema_view,
-            ) = ProjectBuilder._build_label_schema(  # type: ignore[assignment]
-                project=project,
-                labels_by_task=labels_by_task,
-                group_name_by_label_id=group_name_by_label,
-                child_to_parent_id=child_to_parent_id,  # type: ignore
-                previous_schema_revision_id=label_schema.id_,
-                previous_task_id_to_schema_revision_id=task_id_to_schema_revision_id,
-            )
-            label_schema_repo.save(instance=label_schema)
-            for _, task_label_schema in task_to_label_schema_view.items():
-                label_schema_repo.save(instance=task_label_schema)
+        (
+            new_label_schema,
+            new_task_to_label_schema_view,
+        ) = ProjectBuilder._build_label_schema(  # type: ignore[assignment]
+            project=project,
+            labels_by_task=labels_by_task,
+            group_name_by_label_id=group_name_by_label,
+            child_to_parent_id=child_to_parent_id,  # type: ignore
+            previous_schema_revision_id=label_schema.id_,
+            previous_task_id_to_schema_revision_id=task_id_to_schema_revision_id,
+        )
+
+        # Reset label schema (view) ids if no structural change was made to the label topology
+        # Note that we always have to regenerate and save the schemas to ensure the order of the labels is updated
+        if not labels_structure_changed:
+            new_label_schema.id_ = label_schema.id_
+            new_label_schema.previous_schema_revision_id = label_schema.previous_schema_revision_id
+            for task_title, new_task_label_schema_view in new_task_to_label_schema_view.items():
+                new_task_label_schema_view.id_ = task_to_label_schema_view[task_title].id_
+                new_task_label_schema_view.previous_schema_revision_id = task_to_label_schema_view[
+                    task_title
+                ].previous_schema_revision_id
+
+        label_schema_repo.save(instance=new_label_schema)
+        for new_task_label_schema in new_task_to_label_schema_view.values():
+            label_schema_repo.save(instance=new_task_label_schema)
 
         project_repo.save(instance=project)
         return (
             project,
-            label_schema,
-            task_to_label_schema_view,
+            new_label_schema,
+            new_task_to_label_schema_view,
             tuple(labels_to_revisit),
             modified_scene_ids_by_storage,
             labels_structure_changed,

interactive_ai/libs/iai_core_py/tests/entities/test_graph.py

@@ -851,19 +851,3 @@ def test_multi_di_graph(self):
                 },
             ]
         )
-
-    def test_multi_di_graph_topological_sort(self):
-        """
-        <b>Description:</b>
-        Check topological_sort method of MultiDiGraph class object
-
-        <b>Input data:</b>
-        MultiDiGraph objects with specified "edges" parameter
-
-        <b>Expected results:</b>
-        Test passes if topological_sort method returns generator object with expected values
-        """
-        multi_di_graph = self.multi_di_graph()
-        topological_sort = multi_di_graph.topological_sort()
-        for expected_value in [(1, 1), (1, 2), (3, 1), (2, 1)]:
-            assert next(topological_sort) == expected_value