diff --git a/libs/labelbox/src/labelbox/data/annotation_types/relationship.py b/libs/labelbox/src/labelbox/data/annotation_types/relationship.py index a5b84b3ba..29c47813c 100644 --- a/libs/labelbox/src/labelbox/data/annotation_types/relationship.py +++ b/libs/labelbox/src/labelbox/data/annotation_types/relationship.py @@ -14,7 +14,8 @@ class Type(Enum): UNIDIRECTIONAL = "unidirectional" BIDIRECTIONAL = "bidirectional" - source: Union[ObjectAnnotation, ClassificationAnnotation] + source: Optional[Union[ObjectAnnotation, ClassificationAnnotation]] = None + source_ontology_name: Optional[str] = None target: ObjectAnnotation type: Type = Type.UNIDIRECTIONAL readonly: Optional[bool] = None @@ -27,6 +28,22 @@ def check_readonly(self): ) return self + @model_validator(mode="after") + def validate_source_fields(self): + if self.source is None and self.source_ontology_name is None: + raise ValueError( + "Either source or source_ontology_name must be provided" + ) + return self + + @model_validator(mode="after") + def validate_source_consistency(self): + if self.source is not None and self.source_ontology_name is not None: + raise ValueError( + "Only one of 'source' or 'source_ontology_name' may be provided" + ) + return self + class RelationshipAnnotation(BaseAnnotation): value: Relationship diff --git a/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py b/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py index 5b146b660..2f4799d13 100644 --- a/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py +++ b/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py @@ -209,6 +209,8 @@ def _create_relationship_annotations( - For other target annotations: source must be ObjectAnnotation - Target: - Target must always be ObjectAnnotation + ValueError: If relationship validation fails: + - For PDF target annotations: either source or source_ontology_name must be provided """ for annotation in label.annotations: if isinstance(annotation, RelationshipAnnotation): @@ -221,13 +223,20 @@ def _create_relationship_annotations( if isinstance( target.value, (DocumentRectangle, DocumentEntity) ): - if not isinstance( + if source is not None and 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)}" ) + if ( + source is None + and annotation.value.source_ontology_name is None + ): + raise ValueError( + "Unable to create relationship - either source or source_ontology_name must be provided" + ) elif not isinstance(source, ObjectAnnotation): raise TypeError( f"Unable to create relationship with non ObjectAnnotation source: {type(source)}" @@ -239,7 +248,7 @@ def _create_relationship_annotations( f"Unable to create relationship with non ObjectAnnotation target: {type(target)}" ) - if not source._uuid: + if source is not None and not source._uuid: source._uuid = uuid1 if not target._uuid: target._uuid = uuid2 diff --git a/libs/labelbox/src/labelbox/data/serialization/ndjson/relationship.py b/libs/labelbox/src/labelbox/data/serialization/ndjson/relationship.py index 8495fdc04..1ec9b8d4a 100644 --- a/libs/labelbox/src/labelbox/data/serialization/ndjson/relationship.py +++ b/libs/labelbox/src/labelbox/data/serialization/ndjson/relationship.py @@ -11,10 +11,11 @@ class _Relationship(BaseModel): - source: str + source: Optional[str] = None target: str type: str readonly: Optional[bool] = None + sourceOntologyName: Optional[str] = None class NDRelationship(NDAnnotation): @@ -25,11 +26,13 @@ def to_common( annotation: "NDRelationship", source: SUPPORTED_ANNOTATIONS, target: SUPPORTED_ANNOTATIONS, + source_ontology_name: Optional[str] = None, ) -> RelationshipAnnotation: return RelationshipAnnotation( name=annotation.name, value=Relationship( source=source, + source_ontology_name=source_ontology_name, target=target, type=Relationship.Type(annotation.relationship.type), readonly=annotation.relationship.readonly, @@ -50,8 +53,11 @@ def from_common( name=annotation.name, dataRow=DataRow(id=data.uid, global_key=data.global_key), relationship=_Relationship( - source=str(relationship.source._uuid), + source=str(relationship.source._uuid) + if relationship.source + else None, target=str(relationship.target._uuid), + sourceOntologyName=relationship.source_ontology_name, type=relationship.type.value, readonly=relationship.readonly, ), diff --git a/libs/labelbox/tests/data/annotation_import/test_relationships.py b/libs/labelbox/tests/data/annotation_import/test_relationships.py index 695bac9b2..68d3e538a 100644 --- a/libs/labelbox/tests/data/annotation_import/test_relationships.py +++ b/libs/labelbox/tests/data/annotation_import/test_relationships.py @@ -409,3 +409,68 @@ def test_relationship_readonly_explicit_true(): ), ) assert relationship.value.readonly is True + + +def test_relationship_source_ontology_name(): + """Test that relationship can be created with source_ontology_name instead of source.""" + target = ObjectAnnotation( + name="e2", + value=TextEntity(start=30, end=35), + ) + + relationship = RelationshipAnnotation( + name="rel", + value=Relationship( + source_ontology_name="test_source", + target=target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) + assert relationship.value.source_ontology_name == "test_source" + assert relationship.value.source is None + + +def test_relationship_missing_source_validation(): + """Test that relationship requires either source or source_ontology_name.""" + target = ObjectAnnotation( + name="e2", + value=TextEntity(start=30, end=35), + ) + + with pytest.raises( + ValueError, + match="Either source or source_ontology_name must be provided", + ): + RelationshipAnnotation( + name="rel", + value=Relationship( + target=target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) + + +def test_relationship_both_sources_validation(): + """Test that relationship cannot have both source and source_ontology_name.""" + source = ObjectAnnotation( + name="e1", + value=TextEntity(start=10, end=12), + ) + target = ObjectAnnotation( + name="e2", + value=TextEntity(start=30, end=35), + ) + + with pytest.raises( + ValueError, + match="Only one of 'source' or 'source_ontology_name' may be provided", + ): + RelationshipAnnotation( + name="rel", + value=Relationship( + source=source, + source_ontology_name="test_source", + target=target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) diff --git a/libs/labelbox/tests/data/serialization/ndjson/test_relationship.py b/libs/labelbox/tests/data/serialization/ndjson/test_relationship.py index 7aea28d95..9d3aa6178 100644 --- a/libs/labelbox/tests/data/serialization/ndjson/test_relationship.py +++ b/libs/labelbox/tests/data/serialization/ndjson/test_relationship.py @@ -5,6 +5,9 @@ RelationshipAnnotation, Relationship, TextEntity, + DocumentRectangle, + RectangleUnit, + Point, ) @@ -285,3 +288,140 @@ def test_readonly_relationships(): non_readonly_rel_serialized["relationship"]["type"] == "bidirectional" ) assert non_readonly_rel_serialized["relationship"]["readonly"] is False + + +def test_source_ontology_name_relationship(): + ner_source = ObjectAnnotation( + name="e1", + value=TextEntity(start=10, end=12), + ) + + ner_target = ObjectAnnotation( + name="e2", + value=TextEntity(start=30, end=35), + ) + + # Test relationship with source + regular_relationship = RelationshipAnnotation( + name="regular_rel", + value=Relationship( + source=ner_source, + target=ner_target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) + + # Test relationship with source_ontology_name for PDF target + pdf_target = ObjectAnnotation( + name="pdf_region", + value=DocumentRectangle( + start=Point(x=0.5, y=0.5), + end=Point(x=0.7, y=0.7), + page=1, + unit=RectangleUnit.PERCENT, + ), + ) + + ontology_relationship = RelationshipAnnotation( + name="ontology_rel", + value=Relationship( + source_ontology_name="Person", + target=pdf_target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) + + label = Label( + data={"uid": "clqbkpy236syk07978v3pscw1"}, + annotations=[ + ner_source, + ner_target, + pdf_target, + regular_relationship, + ontology_relationship, + ], + ) + + serialized_label = list(NDJsonConverter.serialize([label])) + + ner_source_serialized = next( + annotation + for annotation in serialized_label + if annotation["name"] == ner_source.name + ) + ner_target_serialized = next( + annotation + for annotation in serialized_label + if annotation["name"] == ner_target.name + ) + pdf_target_serialized = next( + annotation + for annotation in serialized_label + if annotation["name"] == pdf_target.name + ) + regular_rel_serialized = next( + annotation + for annotation in serialized_label + if annotation["name"] == regular_relationship.name + ) + ontology_rel_serialized = next( + annotation + for annotation in serialized_label + if annotation["name"] == ontology_relationship.name + ) + + # Verify regular relationship + assert ( + regular_rel_serialized["relationship"]["source"] + == ner_source_serialized["uuid"] + ) + assert ( + regular_rel_serialized["relationship"]["target"] + == ner_target_serialized["uuid"] + ) + assert regular_rel_serialized["relationship"]["type"] == "unidirectional" + + # Verify relationship with source_ontology_name + assert ( + ontology_rel_serialized["relationship"]["sourceOntologyName"] + == "Person" + ) + assert ( + ontology_rel_serialized["relationship"]["target"] + == pdf_target_serialized["uuid"] + ) + assert ontology_rel_serialized["relationship"]["type"] == "unidirectional" + + # Test that providing both source and source_ontology_name raises an error + try: + RelationshipAnnotation( + name="invalid_rel", + value=Relationship( + source=ner_source, + source_ontology_name="Person", + target=pdf_target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) + assert False, "Expected ValueError for providing both source and source_ontology_name" + except Exception as e: + assert ( + "Value error, Only one of 'source' or 'source_ontology_name' may be provided" + in str(e) + ) + + # Test that providing neither source nor source_ontology_name raises an error + try: + RelationshipAnnotation( + name="invalid_rel", + value=Relationship( + target=pdf_target, + type=Relationship.Type.UNIDIRECTIONAL, + ), + ) + assert False, "Expected ValueError for providing neither source nor source_ontology_name" + except Exception as e: + assert ( + "Value error, Either source or source_ontology_name must be provided" + in str(e) + )