Fix 'None' node issue in label schema mapping in case of label deletion (#2300)

Signed-off-by: Songki Choi <[email protected]>

Author: Songki Choi

otx/api/serialization/label_mapper.py

@@ -142,9 +142,13 @@ def backward(instance: dict, all_labels: Dict[ID, LabelEntity]) -> Union[LabelTr
 
         label_map = {label_id: all_labels.get(IDMapper().backward(label_id)) for label_id in instance["nodes"]}
         for label in label_map.values():
-            output.add_node(label)
+            if label:
+                output.add_node(label)
         for edge in instance["edges"]:
-            output.add_edge(label_map[edge[0]], label_map[edge[1]])
+            node1 = label_map.get(edge[0])
+            node2 = label_map.get(edge[1])
+            if node1 and node2:
+                output.add_edge(node1, node2)
 
         return output
 

tests/unit/api/serialization/test_label_mapper.py

@@ -374,3 +374,30 @@ def test_label_graph_backward(self):
         }
         with pytest.raises(ValueError):
             LabelGraphMapper.backward(instance=forward, all_labels=labels)
+        # Checking label deletion case
+        forward = {
+            "type": "tree",
+            "directed": True,
+            "nodes": ["0", "0_1", "0_2", "0_1_1", "0_2_1"],
+            "edges": [("0_1", "0"), ("0_2", "0"), ("0_1_1", "0_1"), ("0_2_1", "0_2")],
+        }
+        labels = {
+            ID("0"): self.label_0,
+            ID("0_1"): self.label_0_1,
+            # ID("0_2"): self.label_0_2,
+            ID("0_1_1"): self.label_0_1_1,
+            # ID("0_1_2"): self.label_0_1_2,
+            ID("0_2_1"): self.label_0_2_1,
+        }
+        expected_backward = LabelTree()
+        for node in labels.values():
+            expected_backward.add_node(node)
+        for parent, child in [
+            (self.label_0, self.label_0_1),
+            # (self.label_0, self.label_0_2),
+            (self.label_0_1, self.label_0_1_1),
+            # (self.label_0_2, self.label_0_2_1),
+        ]:
+            expected_backward.add_child(parent, child)
+        actual_backward = LabelGraphMapper.backward(instance=forward, all_labels=labels)
+        assert LabelGraphMapper.forward(actual_backward) == LabelGraphMapper.forward(expected_backward)