support new api types

Author: jihan-yin

nucleus/constants.py

@@ -11,6 +11,7 @@
 ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE, CUBOID_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
+AUTOTAG_SCORE_THRESHOLD = "score_threshold"
 EXPORTED_ROWS = "exportedRows"
 CAMERA_PARAMS_KEY = "camera_params"
 CLASS_PDF_KEY = "class_pdf"
@@ -62,7 +63,8 @@
 MODEL_RUN_ID_KEY = "model_run_id"
 NAME_KEY = "name"
 NEW_ITEMS = "new_items"
-NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
+NUCLEUS_ENDPOINT = "https://jihan-yin-api.devbox.internal.scale.com/v1/nucleus"
+# NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 NUM_SENSORS_KEY = "num_sensors"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
 POINTCLOUD_KEY = "pointcloud"

nucleus/dataset.py

@@ -24,6 +24,7 @@
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    AUTOTAG_SCORE_THRESHOLD,
     UPDATE_KEY,
 )
 from .dataset_item import (
@@ -86,21 +87,36 @@ def items(self) -> List[DatasetItem]:
         return self._client.get_dataset_items(self.id)
 
     @sanitize_string_args
-    def autotag_scores(self, autotag_name, for_scores_greater_than=0):
-        """Export the autotag scores above a threshold, largest scores first.
+    def autotag_items(self, autotag_name, for_scores_greater_than=0):
+        """For a given Autotag of this dataset, export its tagged items with scores above a threshold, largest scores first.
 
-        If you have pandas installed, you can create a pandas dataframe using
+        :return: dictionary of the form
+            {
+                'autotagItems': AutotagItem[],
+                'autotag': Autotag
+            }
+        See https://dashboard.nucleus.scale.com/nucleus/docs/api#export-autotag-items for more details on the return types.
+        """
+        response = self._client.make_request(
+            payload={AUTOTAG_SCORE_THRESHOLD: for_scores_greater_than},
+            route=f"autotag/dataset/{self.id}/autotag/{autotag_name}/taggedItems",
+            requests_command=requests.get,
+        )
+        return response
 
-        pandas.Dataframe(dataset.autotag_scores(autotag_name))
+    def autotag_training_items(self, autotag_name):
+        """For a given Autotag of this dataset, export its training items. These are user selected positives during refinement.
 
         :return: dictionary of the form
-            {'ref_ids': List[str],
-             'datset_item_ids': List[str],
-             'score': List[float]}
+            {
+                'autotagTrainingItems': AutotagTrainingItem[],
+                'autotag': Autotag
+            }
+        See https://dashboard.nucleus.scale.com/nucleus/docs/api#export-autotag-training-items for more details on the return types.
         """
         response = self._client.make_request(
             payload={},
-            route=f"autotag/{self.id}/{autotag_name}/{for_scores_greater_than}",
+            route=f"autotag/dataset/{self.id}/autotag/{autotag_name}/trainingItems",
             requests_command=requests.get,
         )
         return response