diff --git a/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py b/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py index 9ceb5dafd..e822f3c42 100644 --- a/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py +++ b/libs/labelbox/src/labelbox/data/serialization/ndjson/label.py @@ -65,11 +65,7 @@ def from_common( cls, data: LabelCollection ) -> Generator["NDLabel", None, None]: for label in data: - if all( - isinstance(model, RelationshipAnnotation) - for model in label.annotations - ): - yield from cls._create_relationship_annotations(label) + yield from cls._create_relationship_annotations(label) yield from cls._create_non_video_annotations(label) yield from cls._create_video_annotations(label) @@ -194,25 +190,22 @@ def _create_non_video_annotations(cls, label: Label): f"Unable to convert object to MAL format. `{type(getattr(annotation, 'value',annotation))}`" ) + @classmethod def _create_relationship_annotations(cls, label: Label): - relationship_annotations = [ - annotation - for annotation in label.annotations - if isinstance(annotation, RelationshipAnnotation) - ] - for relationship_annotation in relationship_annotations: - uuid1 = uuid4() - uuid2 = uuid4() - source = copy.copy(relationship_annotation.value.source) - target = copy.copy(relationship_annotation.value.target) - if not isinstance(source, ObjectAnnotation) or not isinstance( - target, ObjectAnnotation - ): - raise TypeError( - f"Unable to create relationship with non ObjectAnnotations. `Source: {type(source)} Target: {type(target)}`" - ) - if not source._uuid: - source._uuid = uuid1 - if not target._uuid: - target._uuid = uuid2 - yield relationship_annotation + 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 + ): + raise TypeError( + f"Unable to create relationship with non ObjectAnnotations. `Source: {type(source)} Target: {type(target)}`" + ) + if not source._uuid: + source._uuid = uuid1 + if not target._uuid: + target._uuid = uuid2 + yield NDRelationship.from_common(annotation, label.data) diff --git a/libs/labelbox/tests/data/annotation_import/test_relationships.py b/libs/labelbox/tests/data/annotation_import/test_relationships.py index 038767e49..1335261e5 100644 --- a/libs/labelbox/tests/data/annotation_import/test_relationships.py +++ b/libs/labelbox/tests/data/annotation_import/test_relationships.py @@ -145,6 +145,7 @@ def configured_project( rand_gen, data_row_json_by_media_type, normalized_ontology_by_media_type_relationship, + teardown_helpers, ): """Configure project for test. Request.param will contain the media type if not present will use Image MediaType. The project will have 10 data rows.""" @@ -186,6 +187,10 @@ def configured_project( project.global_keys = global_keys yield project + teardown_helpers.teardown_project_labels_ontology_feature_schemas(project) + + if dataset: + dataset.delete() @pytest.mark.parametrize(