[SDK-0] Fix class annotations (#1349)

Author: mnoszczak

labelbox/data/annotation_types/base_annotation.py

@@ -1,6 +1,6 @@
 import abc
-from uuid import UUID, uuid4
-from typing import Any, Dict
+from uuid import UUID
+from typing import Any, Dict, Optional
 from pydantic import PrivateAttr
 
 from .feature import FeatureSchema
@@ -9,10 +9,10 @@
 class BaseAnnotation(FeatureSchema, abc.ABC):
     """ Base annotation class. Shouldn't be directly instantiated
     """
-    _uuid: UUID = PrivateAttr()
+    _uuid: Optional[UUID] = PrivateAttr()
     extra: Dict[str, Any] = {}
 
     def __init__(self, **data):
         super().__init__(**data)
         extra_uuid = data.get("extra", {}).get("uuid")
-        self._uuid = data.get("_uuid") or extra_uuid or uuid4()
+        self._uuid = data.get("_uuid") or extra_uuid or None

labelbox/data/serialization/ndjson/converter.py

@@ -1,7 +1,9 @@
 import logging
+import uuid
 from typing import Any, Dict, Generator, Iterable
 
 from ...annotation_types.collection import LabelCollection, LabelGenerator
+from ...annotation_types.relationship import RelationshipAnnotation
 from .label import NDLabel
 
 logger = logging.getLogger(__name__)
@@ -21,7 +23,7 @@ def deserialize(json_data: Iterable[Dict[str, Any]]) -> LabelGenerator:
         Returns:
             LabelGenerator containing the ndjson data.
         """
-        data = NDLabel(**{'annotations': json_data})
+        data = NDLabel(**{"annotations": json_data})
         res = data.to_common()
         return res
 
@@ -40,8 +42,50 @@ def serialize(
         Returns:
             A generator for accessing the ndjson representation of the data
         """
+        used_annotation_uuids = set()
+        for label in labels:
+            annotation_uuid_to_generated_uuid_lookup = {}
+            # UUIDs are private properties used to enhance UX when defining relationships.
+            # They are created for all annotations, but only utilized for relationships.
+            # To avoid overwriting, UUIDs must be unique across labels.
+            # Non-relationship annotation UUIDs are dropped (server-side generation will occur).
+            # For relationship annotations, new UUIDs are generated and stored in a lookup table.
+            for annotation in label.annotations:
+                if isinstance(annotation, RelationshipAnnotation):
+                    source_uuid = annotation.value.source._uuid
+                    target_uuid = annotation.value.target._uuid
+
+                    if (len(
+                            used_annotation_uuids.intersection(
+                                {source_uuid, target_uuid})) > 0):
+                        new_source_uuid = uuid.uuid4()
+                        new_target_uuid = uuid.uuid4()
+
+                        annotation_uuid_to_generated_uuid_lookup[
+                            source_uuid] = new_source_uuid
+                        annotation_uuid_to_generated_uuid_lookup[
+                            target_uuid] = new_target_uuid
+                        annotation.value.source._uuid = new_source_uuid
+                        annotation.value.target._uuid = new_target_uuid
+                    else:
+                        annotation_uuid_to_generated_uuid_lookup[
+                            source_uuid] = source_uuid
+                        annotation_uuid_to_generated_uuid_lookup[
+                            target_uuid] = target_uuid
+                    used_annotation_uuids.add(annotation._uuid)
+
+            for annotation in label.annotations:
+                if (not isinstance(annotation, RelationshipAnnotation) and
+                        hasattr(annotation, "_uuid")):
+                    annotation._uuid = annotation_uuid_to_generated_uuid_lookup.get(
+                        annotation._uuid, annotation._uuid)
+
         for example in NDLabel.from_common(labels):
-            res = example.dict(by_alias=True)
+            annotation_uuid = getattr(example, "uuid", None)
+
+            res = example.dict(
+                by_alias=True,
+                exclude={"uuid"} if annotation_uuid == "None" else None)
             for k, v in list(res.items()):
                 if k in IGNORE_IF_NONE and v is None:
                     del res[k]

tests/data/serialization/ndjson/test_checklist.py

@@ -33,7 +33,6 @@ def test_serialization_min():
     serialized = NDJsonConverter.serialize([label])
     res = next(serialized)
 
-    res.pop("uuid")
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize([res])
@@ -113,7 +112,6 @@ def test_serialization_with_classification():
     serialized = NDJsonConverter.serialize([label])
     res = next(serialized)
 
-    res.pop("uuid")
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize([res])
@@ -197,7 +195,6 @@ def test_serialization_with_classification_double_nested():
     serialized = NDJsonConverter.serialize([label])
     res = next(serialized)
 
-    res.pop("uuid")
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize([res])
@@ -277,7 +274,6 @@ def test_serialization_with_classification_double_nested_2():
 
     serialized = NDJsonConverter.serialize([label])
     res = next(serialized)
-    res.pop("uuid")
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize([res])

tests/data/serialization/ndjson/test_conversation.py

@@ -83,7 +83,6 @@
      [free_text_label, free_text_ndjson]])
 def test_message_based_radio_classification(label, ndjson):
     serialized_label = list(NDJsonConverter().serialize(label))
-    serialized_label[0].pop('uuid')
     assert serialized_label == ndjson
 
     deserialized_label = list(NDJsonConverter().deserialize(ndjson))

tests/data/serialization/ndjson/test_document.py

@@ -69,7 +69,6 @@ def test_pdf_with_name_only():
 
 def test_pdf_bbox_serialize():
     serialized = list(NDJsonConverter.serialize(bbox_labels))
-    serialized[0].pop('uuid')
     assert serialized == bbox_ndjson
 
 

tests/data/serialization/ndjson/test_image.py

@@ -85,7 +85,6 @@ def test_mask_from_arr():
     ],
                   data=ImageData(uid="0" * 25))
     res = next(NDJsonConverter.serialize([label]))
-    res.pop('uuid')
     assert res == {
         "classifications": [],
         "schemaId": "1" * 25,

tests/data/serialization/ndjson/test_radio.py

@@ -34,7 +34,6 @@ def test_serialization_with_radio_min():
     serialized = NDJsonConverter.serialize([label])
     res = next(serialized)
 
-    res.pop("uuid")
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize([res])
@@ -86,7 +85,6 @@ def test_serialization_with_radio_classification():
 
     serialized = NDJsonConverter.serialize([label])
     res = next(serialized)
-    res.pop("uuid")
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize([res])

tests/data/serialization/ndjson/test_video.py

@@ -87,8 +87,6 @@ def test_video_classification_global_subclassifications():
 
     serialized = NDJsonConverter.serialize([label])
     res = [x for x in serialized]
-    for annotations in res:
-        annotations.pop("uuid")
     assert res == [expected_first_annotation, expected_second_annotation]
 
     deserialized = NDJsonConverter.deserialize(res)

tests/integration/annotation_import/test_data_types.py

@@ -59,7 +59,7 @@
 def get_annotation_comparison_dicts_from_labels(labels):
     labels_ndjson = list(NDJsonConverter.serialize(labels))
     for annotation in labels_ndjson:
-        annotation.pop('uuid')
+        annotation.pop('uuid', None)
         annotation.pop('dataRow')
 
         if 'masks' in annotation: