Merge pull request #31 from scaleapi/sasha/list_autotags

Sasha/list autotags

Author: sasha-scale

nucleus/__init__.py

@@ -103,6 +103,7 @@
     DATASET_LENGTH_KEY,
     NAME_KEY,
     ANNOTATIONS_KEY,
+    AUTOTAGS_KEY,
 )
 from .model import Model
 from .errors import (
@@ -838,6 +839,19 @@ def delete_slice(self, slice_id: str) -> dict:
         )
         return response
 
+    def list_autotags(self, dataset_id: str) -> List[str]:
+        """
+        Fetches a list of autotags for a given dataset id
+        :param dataset_id: internally controlled dataset_id
+        :return: List[str] representing autotag_ids
+        """
+        response = self._make_request(
+            {},
+            f"{dataset_id}/list_autotags",
+            requests_command=requests.get,
+        )
+        return response[AUTOTAGS_KEY] if AUTOTAGS_KEY in response else response
+
     def delete_model(self, model_id: str) -> dict:
         """
         This endpoint deletes the specified model, along with all

nucleus/constants.py

@@ -42,3 +42,4 @@
 BOX_TYPE = "box"
 POLYGON_TYPE = "polygon"
 GEOMETRY_KEY = "geometry"
+AUTOTAGS_KEY = "autotags"

nucleus/dataset.py

@@ -221,6 +221,9 @@ def delete_item(self, item_id: str = None, reference_id: str = None):
             self.id, reference_id=reference_id, item_id=item_id
         )
 
+    def list_autotags(self):
+        return self._client.list_autotags(self.id)
+
     def _format_dataset_item_response(self, response: dict) -> dict:
         item = response.get(ITEM_KEY, None)
         annotation_payload = response.get(ANNOTATIONS_KEY, [])

nucleus/slice.py

@@ -1,5 +1,9 @@
 from typing import List, Dict, Any
-from .constants import DATASET_ITEM_ID_KEY, REFERENCE_IDS_KEY
+from .constants import (
+    DATASET_ITEM_ID_KEY,
+    REFERENCE_IDS_KEY,
+    DATASET_ITEM_IDS_KEY,
+)
 
 
 class Slice:
@@ -29,9 +33,9 @@ def info(self, id_type: str = DATASET_ITEM_ID_KEY) -> dict:
         return self._client.slice_info(self.slice_id, id_type)
 
     def append(
-        self, 
+        self,
         dataset_item_ids: List[str] = None,
-        reference_ids: List[str] = None
+        reference_ids: List[str] = None,
     ) -> dict:
         """
         Appends to a slice from items already present in a dataset.
@@ -61,4 +65,4 @@ def append(
         response = self._client._make_request(
             payload, f"slice/{self.slice_id}/append"
         )
-        return response
+        return response

tests/test_dataset.py

@@ -1,7 +1,5 @@
-import pytest
-
 from pathlib import Path
-
+import pytest
 from nucleus import Dataset, DatasetItem, UploadResponse
 from nucleus.constants import (
     NEW_ITEMS,
@@ -12,15 +10,14 @@
     DATASET_ID_KEY,
 )
 
-
 TEST_DATASET_NAME = '[PyTest] Test Dataset'
 TEST_SLICE_NAME = '[PyTest] Test Slice'
 TEST_IMG_URLS = [
-    's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg',
-    's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg',
-    's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/7f2e1814-6591087d.jpg',
-    's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg',
-    's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg',
+    "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg",
+    "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg",
+    "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/7f2e1814-6591087d.jpg",
+    "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg",
+    "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg",
 ]
 
 @pytest.fixture()
@@ -66,25 +63,33 @@ def check_is_expected_response(response):
     # With reference ids and metadata:
     ds_items_with_metadata = []
     for i, url in enumerate(TEST_IMG_URLS):
-        ds_items_with_metadata.append(DatasetItem(
-            image_location=url,
-            reference_id=Path(url).name,
-            metadata={
-                'made_with_pytest': True,
-                'example_int': i,
-                'example_str': 'hello',
-                'example_float': 0.5,
-                'example_dict': {
-                    'nested': True,
+        ds_items_with_metadata.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=Path(url).name,
+                metadata={
+                    "made_with_pytest": True,
+                    "example_int": i,
+                    "example_str": "hello",
+                    "example_float": 0.5,
+                    "example_dict": {
+                        "nested": True,
+                    },
+                    "example_list": ["hello", i, False],
                 },
-                'example_list': ['hello', i, False],
-            }
-        ))
+            )
+        )
 
     response = dataset.append(ds_items_with_metadata)
     check_is_expected_response(response)
 
 
+def test_dataset_list_autotags(CLIENT, dataset):
+    # Creation
+    # List of Autotags should be empty
+    autotag_response = CLIENT.list_autotags(dataset.id)
+    assert autotag_response == []
+
 def test_slice_create_and_delete(dataset):
     # Dataset upload
     ds_items = []