add integration test for 3d dataset creation and cuboid upload

Author: Drew Kaul

tests/helpers.py

@@ -9,6 +9,7 @@
 TEST_MODEL_NAME = "[PyTest] Test Model"
 TEST_MODEL_RUN = "[PyTest] Test Model Run"
 TEST_DATASET_NAME = "[PyTest] Test Dataset"
+TEST_DATASET_3D_NAME = "[PyTest] Test Dataset 3D"
 TEST_SLICE_NAME = "[PyTest] Test Slice"
 TEST_PROJECT_ID = "60b699d70f139e002dd31bfc"
 
@@ -21,6 +22,10 @@
     "https://homepages.cae.wisc.edu/~ece533/images/cat.png",
 ]
 
+TEST_POINTCLOUD_URLS = [
+    "https://scale-us-attachments.s3.us-west-2.amazonaws.com/select/examples/kitti_example_frame.json",
+]
+
 TEST_DATASET_ITEMS = [
     DatasetItem(TEST_IMG_URLS[0], "1"),
     DatasetItem(TEST_IMG_URLS[1], "2"),
@@ -72,6 +77,28 @@ def reference_id_from_url(url):
     for i in range(len(TEST_IMG_URLS))
 ]
 
+TEST_CUBOID_ANNOTATIONS = [
+    {
+        "label": f"[Pytest] Cuboid Annotation ${i}",
+        "geometry": {
+            "position": {
+                "x": 50 * i + 10,
+                "y": 60 * i + 10,
+                "z": 70 * i + 10,
+            },
+            "dimensions": {
+                "x": 10 * i,
+                "y": 20 * i,
+                "z": 30 * i,
+            },
+            "yaw": 5 * i,
+        },
+        "reference_id": reference_id_from_url(TEST_POINTCLOUD_URLS[i]),
+        "annotation_id": f"[Pytest] Cuboid Annotation Annotation Id{i}",
+    }
+    for i in range(len(TEST_POINTCLOUD_URLS))
+]
+
 
 TEST_MASK_URL = "https://raw.githubusercontent.com/scaleapi/nucleus-python-client/master/tests/testdata/000000000285.png"
 
@@ -157,6 +184,29 @@ def assert_polygon_annotation_matches_dict(
         assert instance_pt.y == dict_pt["y"]
 
 
+def assert_cuboid_annotation_matches_dict(
+    annotation_instance, annotation_dict
+):
+    assert annotation_instance.label == annotation_dict["label"]
+    assert (
+        annotation_instance.annotation_id == annotation_dict["annotation_id"]
+    )
+    for instance_pt, dict_pt in zip(
+        annotation_instance.position, annotation_dict["geometry"]["position"]
+    ):
+        assert instance_pt.x == dict_pt["x"]
+        assert instance_pt.y == dict_pt["y"]
+        assert instance_pt.z == dict_pt["z"]
+    for instance_pt, dict_pt in zip(
+        annotation_instance.dimensions,
+        annotation_dict["geometry"]["dimensions"],
+    ):
+        assert instance_pt.x == dict_pt["x"]
+        assert instance_pt.y == dict_pt["y"]
+        assert instance_pt.z == dict_pt["z"]
+    assert annotation_instance.yaw == annotation_dict["geometry"]["yaw"]
+
+
 def assert_segmentation_annotation_matches_dict(
     annotation_instance, annotation_dict
 ):

tests/test_annotation.py

@@ -2,19 +2,24 @@
 
 from .helpers import (
     TEST_DATASET_NAME,
+    TEST_DATASET_3D_NAME,
     TEST_IMG_URLS,
+    TEST_POINTCLOUD_URLS,
     TEST_BOX_ANNOTATIONS,
     TEST_POLYGON_ANNOTATIONS,
+    TEST_CUBOID_ANNOTATIONS,
     TEST_SEGMENTATION_ANNOTATIONS,
     reference_id_from_url,
     assert_box_annotation_matches_dict,
     assert_polygon_annotation_matches_dict,
+    assert_cuboid_annotation_matches_dict,
     assert_segmentation_annotation_matches_dict,
 )
 
 from nucleus import (
     BoxAnnotation,
     PolygonAnnotation,
+    CuboidAnnotation,
     SegmentationAnnotation,
     DatasetItem,
     Segment,
@@ -61,6 +66,26 @@ def dataset(CLIENT):
     assert response == {"message": "Beginning dataset deletion..."}
 
 
+@pytest.fixture()
+def dataset_3d(CLIENT):
+    ds = CLIENT.create_dataset(TEST_DATASET_3D_NAME)
+    ds_items = []
+    for url in TEST_POINTCLOUD_URLS:
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
+
+    response = ds.append(ds_items)
+    assert ERROR_PAYLOAD not in response.json()
+    yield ds
+
+    response = CLIENT.delete_dataset(ds.id)
+    assert response == {"message": "Beginning dataset deletion..."}
+
+
 def test_box_gt_upload(dataset):
     annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
     response = dataset.annotate(annotations=[annotation])
@@ -95,6 +120,24 @@ def test_polygon_gt_upload(dataset):
     )
 
 
+def test_cuboid_gt_upload(dataset_3d):
+    annotation = CuboidAnnotation.from_json(TEST_CUBOID_ANNOTATIONS[0])
+    response = dataset_3d.annotate(annotations=[annotation])
+
+    assert response["dataset_id"] == dataset_3d.id
+    assert response["annotations_processed"] == 1
+    assert response["annotations_ignored"] == 0
+
+    response = dataset_3d.refloc(annotation.reference_id)["annotations"][
+        "cuboid"
+    ]
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_cuboid_annotation_matches_dict(
+        response_annotation, TEST_CUBOID_ANNOTATIONS[0]
+    )
+
+
 def test_single_semseg_gt_upload(dataset):
     annotation = SegmentationAnnotation.from_json(
         TEST_SEGMENTATION_ANNOTATIONS[0]