dataset export works

Author: Diego Ardila

nucleus/dataset.py

@@ -4,6 +4,7 @@
 
 from nucleus.job import AsyncJob
 from nucleus.utils import (
+    convert_export_payload,
     format_dataset_item_response,
     serialize_and_write_to_presigned_url,
 )
@@ -327,3 +328,23 @@ def delete_custom_index(self):
 
     def check_index_status(self, job_id: str):
         return self._client.check_index_status(job_id)
+
+    def items_and_annotations(
+        self,
+    ) -> List[Dict[str, Union[DatasetItem, Dict[str, List[Annotation]]]]]:
+        """Returns a list of all DatasetItems and Annotations in this slice.
+
+        Returns:
+            A list, where each item is a dict with two keys representing a row
+            in the dataset.
+            * One value in the dict is the DatasetItem, containing a reference to the
+                item that was annotated.
+            * The other value is a dictionary containing all the annotations for this
+                dataset item, sorted by annotation type.
+        """
+        api_payload = self._client.make_request(
+            payload=None,
+            route=f"dataset/{self.id}/exportForTraining",
+            requests_command=requests.get,
+        )
+        return convert_export_payload(api_payload["exportedRows"])

tests/test_dataset.py

@@ -1,3 +1,4 @@
+import copy
 import math
 import os
 
@@ -21,6 +22,9 @@
     IGNORED_ITEMS,
     NEW_ITEMS,
     UPDATED_ITEMS,
+    ITEM_KEY,
+    ANNOTATIONS_KEY,
+    BOX_TYPE,
 )
 from nucleus.job import AsyncJob, JobError
 
@@ -331,3 +335,29 @@ def test_annotate_async_with_error(dataset: Dataset):
     }
 
     assert "Item with id fake_garbage doesn" in str(job.errors())
+
+
+def test_append_and_export(dataset):
+    # Dataset upload
+    url = TEST_IMG_URLS[0]
+    annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
+
+    ds_items = [
+        DatasetItem(
+            image_location=url,
+            reference_id=reference_id_from_url(url),
+            metadata={"test": "metadata"},
+        ),
+    ]
+    response = dataset.append(ds_items)
+    assert ERROR_PAYLOAD not in response.json()
+
+    dataset.annotate(annotations=[annotation])
+
+    expected_box_annotation = copy.deepcopy(annotation)
+    expected_box_annotation.annotation_id = None
+    expected_box_annotation.metadata = {}
+
+    exported = dataset.items_and_annotations()
+    assert exported[0][ITEM_KEY] == ds_items[0]
+    assert exported[0][ANNOTATIONS_KEY][BOX_TYPE][0] == expected_box_annotation