[AL-5103] Add new data type classes to annotation types

Author: whistler

labelbox/data/annotation_types/__init__.py

@@ -20,6 +20,11 @@
 from .classification import Radio
 from .classification import Text
 
+from .data import AudioData
+from .data import ConversationData
+from .data import DicomData
+from .data import DocumentData
+from .data import HTMLData
 from .data import ImageData
 from .data import MaskData
 from .data import TextData

labelbox/data/annotation_types/data/__init__.py

@@ -1,4 +1,9 @@
+from .audio import AudioData
+from .conversation import ConversationData
+from .dicom import DicomData
+from .document import DocumentData
+from .html import HTMLData
 from .raster import ImageData
 from .raster import MaskData
 from .text import TextData
-from .video import VideoData
+from .video import VideoData

labelbox/data/annotation_types/data/audio.py

@@ -0,0 +1,5 @@
+from .base_data import BaseData
+
+
+class AudioData(BaseData):
+    ...

labelbox/data/annotation_types/data/conversation.py

@@ -0,0 +1,5 @@
+from .base_data import BaseData
+
+
+class ConversationData(BaseData):
+    ...

labelbox/data/annotation_types/data/dicom.py

@@ -0,0 +1,5 @@
+from .base_data import BaseData
+
+
+class DicomData(BaseData):
+    ...

labelbox/data/annotation_types/data/document.py

@@ -0,0 +1,5 @@
+from .base_data import BaseData
+
+
+class DocumentData(BaseData):
+    ...

labelbox/data/annotation_types/data/html.py

@@ -0,0 +1,5 @@
+from .base_data import BaseData
+
+
+class HTMLData(BaseData):
+    ...

tests/integration/annotation_import/conftest.py

@@ -11,6 +11,167 @@
 from labelbox.schema.queue_mode import QueueMode
 
 
+@pytest.fixture()
+def audio_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://storage.googleapis.com/labelbox-datasets/audio-sample-data/sample-audio-1.mp3",
+        "global_key":
+            f"https://storage.googleapis.com/labelbox-datasets/audio-sample-data/sample-audio-1.mp3-{rand_gen(str)}",
+        "media_type":
+            "AUDIO",
+    }
+
+
+@pytest.fixture()
+def conversation_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://storage.googleapis.com/labelbox-datasets/conversational-sample-data/sample-conversation-1.json",
+        "global_key":
+            f"https://storage.googleapis.com/labelbox-datasets/conversational-sample-data/sample-conversation-1.json-{rand_gen(str)}",
+        "media_type":
+            "CONVERSATIONAL",
+    }
+
+
+@pytest.fixture()
+def dicom_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://storage.googleapis.com/labelbox-datasets/dicom-sample-data/sample-dicom-1.dcm",
+        "global_key":
+            f"https://storage.googleapis.com/labelbox-datasets/dicom-sample-data/sample-dicom-1.dcm-{rand_gen(str)}",
+        "media_type":
+            "DICOM",
+    }
+
+
+@pytest.fixture()
+def geospatial_data_row(rand_gen):
+    return {
+        "row_data": {
+            "tile_layer_url":
+                "https://s3-us-west-1.amazonaws.com/lb-tiler-layers/mexico_city/{z}/{x}/{y}.png",
+            "bounds": [[19.405662413477728, -99.21052827588443],
+                       [19.400498983095076, -99.20534818927473]],
+            "min_zoom":
+                12,
+            "max_zoom":
+                20,
+            "epsg":
+                "EPSG4326",
+        },
+        "global_key":
+            f"https://s3-us-west-1.amazonaws.com/lb-tiler-layers/mexico_city/z/x/y.png-{rand_gen(str)}",
+        "media_type":
+            "TMS_GEO",
+    }
+
+
+@pytest.fixture()
+def html_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://storage.googleapis.com/labelbox-datasets/html_sample_data/sample_html_1.html",
+        "global_key":
+            f"https://storage.googleapis.com/labelbox-datasets/html_sample_data/sample_html_1.html-{rand_gen(str)}",
+    }
+
+
+@pytest.fixture()
+def image_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://lb-test-data.s3.us-west-1.amazonaws.com/image-samples/sample-image-1.jpg",
+        "global_key":
+            f"https://lb-test-data.s3.us-west-1.amazonaws.com/image-samples/sample-image-1.jpg-{rand_gen(str)}",
+        "media_type":
+            "IMAGE",
+    }
+
+
+@pytest.fixture()
+def document_data_row(rand_gen):
+    return {
+        "row_data": {
+            "pdf_url":
+                "https://storage.googleapis.com/labelbox-datasets/arxiv-pdf/data/99-word-token-pdfs/0801.3483.pdf",
+            "text_layer_url":
+                "https://storage.googleapis.com/labelbox-datasets/arxiv-pdf/data/99-word-token-pdfs/0801.3483-lb-textlayer.json"
+        },
+        "global_key":
+            f"https://storage.googleapis.com/labelbox-datasets/arxiv-pdf/data/99-word-token-pdfs/0801.3483.pdf-{rand_gen(str)}",
+        "media_type":
+            "PDF",
+    }
+
+
+@pytest.fixture()
+def text_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://lb-test-data.s3.us-west-1.amazonaws.com/text-samples/sample-text-1.txt",
+        "global_key":
+            f"https://lb-test-data.s3.us-west-1.amazonaws.com/text-samples/sample-text-1.txt-{rand_gen(str)}",
+        "media_type":
+            "TEXT",
+    }
+
+
+@pytest.fixture()
+def video_data_row(rand_gen):
+    return {
+        "row_data":
+            "https://storage.googleapis.com/labelbox-datasets/video-sample-data/sample-video-1.mp4",
+        "global_key":
+            f"https://storage.googleapis.com/labelbox-datasets/video-sample-data/sample-video-1.mp4-{rand_gen(str)}",
+        "media_type":
+            "VIDEO",
+    }
+
+
+@pytest.fixture
+def data_row_json_by_data_type(audio_data_row, conversation_data_row,
+                               dicom_data_row, geospatial_data_row,
+                               html_data_row, image_data_row, document_data_row,
+                               text_data_row, video_data_row):
+    return {
+        'audio': audio_data_row,
+        'conversation': conversation_data_row,
+        'dicom': dicom_data_row,
+        'geospatial': geospatial_data_row,
+        'html': html_data_row,
+        'image': image_data_row,
+        'document': document_data_row,
+        'text': text_data_row,
+        'video': video_data_row,
+    }
+
+
+@pytest.fixture
+def annotations_by_data_type(polygon_inference, rectangle_inference,
+                             line_inference, entity_inference,
+                             checklist_inference, text_inference,
+                             video_checklist_inference):
+    return {
+        'audio': [checklist_inference, text_inference],
+        'conversation': [checklist_inference, text_inference, entity_inference],
+        'dicom': [line_inference],
+        'document': [
+            entity_inference, checklist_inference, text_inference,
+            rectangle_inference
+        ],
+        'html': [text_inference, checklist_inference],
+        'image': [
+            polygon_inference, rectangle_inference, line_inference,
+            checklist_inference, text_inference
+        ],
+        'text': [entity_inference, checklist_inference, text_inference],
+        'video': [video_checklist_inference]
+    }
+
+
 @pytest.fixture
 def ontology():
     bbox_tool = {
@@ -168,10 +329,10 @@ def configured_project_pdf(client, ontology, rand_gen, pdf_url):
 
 
 @pytest.fixture
-def dataset_pdf_entity(client, rand_gen, pdf_entity_data_row):
+def dataset_pdf_entity(client, rand_gen, document_data_row):
     dataset = client.create_dataset(name=rand_gen(str))
     data_row_ids = []
-    data_row = dataset.create_data_row(pdf_entity_data_row)
+    data_row = dataset.create_data_row(document_data_row)
     data_row_ids.append(data_row.uid)
     yield dataset, data_row_ids
     dataset.delete()
@@ -298,11 +459,13 @@ def entity_inference(prediction_id_mapping):
 @pytest.fixture
 def segmentation_inference(prediction_id_mapping):
     segmentation = prediction_id_mapping['superpixel'].copy()
-    segmentation.update(
-        {'mask': {
+    segmentation.update({
+        'mask': {
+            # TODO: Use a real URI
             'instanceURI': "sampleuri",
             'colorRGB': [0, 0, 0]
-        }})
+        }
+    })
     del segmentation['tool']
     return segmentation
 

tests/integration/annotation_import/test_data_types.py

@@ -0,0 +1,46 @@
+import pytest
+import labelbox as lb
+import labelbox.types as lb_types
+from labelbox.data.annotation_types.data import AudioData, ConversationData, DicomData, DocumentData, HTMLData, ImageData, TextData
+from labelbox.data.serialization import NDJsonConverter
+from labelbox.schema.annotation_import import AnnotationImportState
+
+
+# TODO: Add VideoData. Currently label import job finishes without errors but project.export_labels() returns empty list.
+@pytest.mark.parametrize('data_type_class', [
+    AudioData, ConversationData, DicomData, DocumentData, HTMLData, ImageData,
+    TextData
+])
+def test_import_data_types(client, configured_project,
+                           data_row_json_by_data_type, annotations_by_data_type,
+                           data_type_class):
+
+    project_id = configured_project.uid
+
+    data_type_string = data_type_class.__name__[:-4].lower()
+    data_row_ndjson = data_row_json_by_data_type[data_type_string]
+    dataset = next(configured_project.datasets())
+    data_row = dataset.create_data_row(data_row_ndjson)
+
+    annotations_ndjson = annotations_by_data_type[data_type_string]
+    annotations_list = [
+        label.annotations
+        for label in NDJsonConverter.deserialize(annotations_ndjson)
+    ]
+    labels = [
+        lb_types.Label(data=data_type_class(uid=data_row.uid),
+                       annotations=annotations)
+        for annotations in annotations_list
+    ]
+
+    label_import = lb.LabelImport.create_from_objects(
+        client, project_id, f'test-import-{data_type_string}', labels)
+    label_import.wait_until_done()
+
+    assert label_import.state == AnnotationImportState.FINISHED
+    assert len(label_import.errors) == 0
+    exported_labels = configured_project.export_labels(download=True)
+    objects = exported_labels[0]['Label']['objects']
+    classifications = exported_labels[0]['Label']['classifications']
+    assert len(objects) + len(classifications) == len(labels)
+    data_row.delete()

tests/integration/conftest.py

@@ -177,23 +177,6 @@ def pdf_url(client):
     return {"row_data": {"pdf_url": pdf_url,}, "global_key": str(uuid.uuid4())}
 
 
-@pytest.fixture(scope="session")
-def pdf_entity_data_row(client):
-    pdf_url = client.upload_file(
-        'tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483.pdf')
-    text_layer_url = client.upload_file(
-        'tests/assets/arxiv-pdf_data_99-word-token-pdfs_0801.3483-lb-textlayer.json'
-    )
-
-    return {
-        "row_data": {
-            "pdf_url": pdf_url,
-            "text_layer_url": text_layer_url
-        },
-        "global_key": str(uuid.uuid4())
-    }
-
-
 @pytest.fixture
 def project(client, rand_gen):
     project = client.create_project(name=rand_gen(str),