MAL and GT support for pdf relationships

Author: lgluszek

libs/labelbox/src/labelbox/data/annotation_types/relationship.py

@@ -1,8 +1,10 @@
+from typing import Union
 from pydantic import BaseModel
 from enum import Enum
 from labelbox.data.annotation_types.annotation import (
     BaseAnnotation,
     ObjectAnnotation,
+    ClassificationAnnotation,
 )
 
 
@@ -11,7 +13,7 @@ class Type(Enum):
         UNIDIRECTIONAL = "unidirectional"
         BIDIRECTIONAL = "bidirectional"
 
-    source: ObjectAnnotation
+    source: Union[ObjectAnnotation, ClassificationAnnotation]
     target: ObjectAnnotation
     type: Type = Type.UNIDIRECTIONAL
 

libs/labelbox/src/labelbox/data/serialization/ndjson/label.py

@@ -2,7 +2,7 @@
 import copy
 from itertools import groupby
 from operator import itemgetter
-from typing import Generator, List, Tuple, Union
+from typing import Generator, List, Tuple, Union, Iterator, Dict
 from uuid import uuid4
 
 from pydantic import BaseModel
@@ -24,6 +24,7 @@
     VideoMaskAnnotation,
     VideoObjectAnnotation,
 )
+from labelbox.types import DocumentRectangle, DocumentEntity
 from .classification import (
     NDChecklistSubclass,
     NDClassification,
@@ -61,9 +62,7 @@ class NDLabel(BaseModel):
     annotations: AnnotationType
 
     @classmethod
-    def from_common(
-        cls, data: LabelCollection
-    ) -> Generator["NDLabel", None, None]:
+    def from_common(cls, data: LabelCollection) -> Generator["NDLabel", None, None]:
         for label in data:
             yield from cls._create_relationship_annotations(label)
             yield from cls._create_non_video_annotations(label)
@@ -127,16 +126,12 @@ def _create_video_annotations(
             if isinstance(
                 annot, (VideoClassificationAnnotation, VideoObjectAnnotation)
             ):
-                video_annotations[annot.feature_schema_id or annot.name].append(
-                    annot
-                )
+                video_annotations[annot.feature_schema_id or annot.name].append(annot)
             elif isinstance(annot, VideoMaskAnnotation):
                 yield NDObject.from_common(annotation=annot, data=label.data)
 
         for annotation_group in video_annotations.values():
-            segment_frame_ranges = cls._get_segment_frame_ranges(
-                annotation_group
-            )
+            segment_frame_ranges = cls._get_segment_frame_ranges(annotation_group)
             if isinstance(annotation_group[0], VideoClassificationAnnotation):
                 annotation = annotation_group[0]
                 frames_data = []
@@ -169,6 +164,7 @@ def _create_non_video_annotations(cls, label: Label):
                     VideoClassificationAnnotation,
                     VideoObjectAnnotation,
                     VideoMaskAnnotation,
+                    RelationshipAnnotation,
                 ),
             )
         ]
@@ -179,8 +175,6 @@ def _create_non_video_annotations(cls, label: Label):
                 yield NDObject.from_common(annotation, label.data)
             elif isinstance(annotation, (ScalarMetric, ConfusionMatrixMetric)):
                 yield NDMetricAnnotation.from_common(annotation, label.data)
-            elif isinstance(annotation, RelationshipAnnotation):
-                yield NDRelationship.from_common(annotation, label.data)
             elif isinstance(annotation, PromptClassificationAnnotation):
                 yield NDPromptClassification.from_common(annotation, label.data)
             elif isinstance(annotation, MessageEvaluationTaskAnnotation):
@@ -191,19 +185,35 @@ def _create_non_video_annotations(cls, label: Label):
                 )
 
     @classmethod
-    def _create_relationship_annotations(cls, label: Label):
+    def _create_relationship_annotations(
+        cls, label: Label
+    ) -> Generator[NDRelationship, None, None]:
         for annotation in label.annotations:
             if isinstance(annotation, RelationshipAnnotation):
                 uuid1 = uuid4()
                 uuid2 = uuid4()
                 source = copy.copy(annotation.value.source)
                 target = copy.copy(annotation.value.target)
-                if not isinstance(source, ObjectAnnotation) or not isinstance(
-                    target, ObjectAnnotation
-                ):
+
+                # Check if source type is valid based on target type
+                if isinstance(target.value, (DocumentRectangle, DocumentEntity)):
+                    if not isinstance(
+                        source, (ObjectAnnotation, ClassificationAnnotation)
+                    ):
+                        raise TypeError(
+                            f"Unable to create relationship with invalid source. For PDF targets, "
+                            f"source must be ObjectAnnotation or ClassificationAnnotation. Got: {type(source)}"
+                        )
+                elif not isinstance(source, ObjectAnnotation):
                     raise TypeError(
-                        f"Unable to create relationship with non ObjectAnnotations. `Source: {type(source)} Target: {type(target)}`"
+                        f"Unable to create relationship with non ObjectAnnotation source: {type(source)}"
                     )
+
+                if not isinstance(target, ObjectAnnotation):
+                    raise TypeError(
+                        f"Unable to create relationship with non ObjectAnnotation target: {type(target)}"
+                    )
+
                 if not source._uuid:
                     source._uuid = uuid1
                 if not target._uuid:

libs/labelbox/tests/data/annotation_import/test_relationships.py

@@ -10,7 +10,17 @@
     RelationshipAnnotation,
     Relationship,
     TextEntity,
+    DocumentRectangle,
+    DocumentEntity,
+    Point,
+    Text,
+    ClassificationAnnotation,
+    DocumentTextSelection,
+    Radio,
+    ClassificationAnswer,
+    Checklist,
 )
+from labelbox.data.serialization.ndjson import NDJsonConverter
 import pytest
 
 
@@ -169,9 +179,7 @@ def configured_project(
     data_row_data = []
 
     for _ in range(3):
-        data_row_data.append(
-            data_row_json_by_media_type[media_type](rand_gen(str))
-        )
+        data_row_data.append(data_row_json_by_media_type[media_type](rand_gen(str)))
 
     task = dataset.create_data_rows(data_row_data)
     task.wait_till_done()
@@ -220,3 +228,100 @@ def test_import_media_types(
 
     assert label_import.state == AnnotationImportState.FINISHED
     assert len(label_import.errors) == 0
+
+
+def test_valid_classification_relationships():
+    def create_pdf_annotation(target_type: str) -> ObjectAnnotation:
+        if target_type == "bbox":
+            return ObjectAnnotation(
+                name="bbox",
+                value=DocumentRectangle(
+                    start=Point(x=0, y=0),
+                    end=Point(x=0.5, y=0.5),
+                    page=1,
+                    unit="PERCENT",
+                ),
+            )
+        elif target_type == "entity":
+            return ObjectAnnotation(
+                name="entity",
+                value=DocumentEntity(
+                    page=1,
+                    textSelections=[
+                        DocumentTextSelection(token_ids=[], group_id="", page=1)
+                    ],
+                ),
+            )
+        raise ValueError(f"Unknown target type: {target_type}")
+
+    def verify_relationship(source: ClassificationAnnotation, target: ObjectAnnotation):
+        relationship = RelationshipAnnotation(
+            name="relationship",
+            value=Relationship(
+                source=source,
+                target=target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+        label = Label(data={"global_key": "global_key"}, annotations=[relationship])
+        result = list(NDJsonConverter.serialize([label]))
+        assert len(result) == 1
+
+    # Test case 1: Text Classification -> DocumentRectangle
+    text_source = ClassificationAnnotation(name="text", value=Text(answer="test"))
+    verify_relationship(text_source, create_pdf_annotation("bbox"))
+
+    # Test case 2: Text Classification -> DocumentEntity
+    verify_relationship(text_source, create_pdf_annotation("entity"))
+
+    # Test case 3: Radio Classification -> DocumentRectangle
+    radio_source = ClassificationAnnotation(
+        name="sub_radio_question",
+        value=Radio(
+            answer=ClassificationAnswer(
+                name="first_sub_radio_answer",
+                classifications=[
+                    ClassificationAnnotation(
+                        name="second_sub_radio_question",
+                        value=Radio(
+                            answer=ClassificationAnswer(name="second_sub_radio_answer")
+                        ),
+                    )
+                ],
+            )
+        ),
+    )
+    verify_relationship(radio_source, create_pdf_annotation("bbox"))
+
+    # Test case 4: Checklist Classification -> DocumentEntity
+    checklist_source = ClassificationAnnotation(
+        name="sub_checklist_question",
+        value=Checklist(
+            answer=[ClassificationAnswer(name="first_sub_checklist_answer")]
+        ),
+    )
+    verify_relationship(checklist_source, create_pdf_annotation("entity"))
+
+
+def test_classification_relationship_restrictions():
+    """Test all relationship validation error messages."""
+    text = ClassificationAnnotation(name="text", value=Text(answer="test"))
+    point = ObjectAnnotation(name="point", value=Point(x=1, y=1))
+
+    # Test case: Classification -> Point (invalid)
+    # Should fail because classifications can only connect to PDF targets
+    relationship = RelationshipAnnotation(
+        name="relationship",
+        value=Relationship(
+            source=text,
+            target=point,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+
+    with pytest.raises(
+        TypeError,
+        match="Unable to create relationship with non ObjectAnnotation source: .*",
+    ):
+        label = Label(data={"global_key": "test_key"}, annotations=[relationship])
+        list(NDJsonConverter.serialize([label]))