Merge pull request #38 from scaleapi/sasha/fix_bugs

Sasha/fix bugs

Author: sasha-scale

nucleus/constants.py

@@ -1,5 +1,4 @@
-# NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
-NUCLEUS_ENDPOINT = "http://localhost:3000/v1/nucleus"
+NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 ITEMS_KEY = "items"
 ITEM_KEY = "item"
 REFERENCE_ID_KEY = "reference_id"

nucleus/model.py

@@ -1,6 +1,10 @@
 from typing import List, Optional, Dict, Union
 from .dataset import Dataset
-from .prediction import BoxPrediction, PolygonPrediction
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .model_run import ModelRun
 from .constants import (
     NAME_KEY,
@@ -25,7 +29,7 @@ def __init__(
         self._client = client
 
     def __repr__(self):
-        return f'Model(model_id={self.id}, name={self.name}, reference_id={self.reference_id}, metadata={self.metadata}, client={self._client})'
+        return f"Model(model_id={self.id}, name={self.name}, reference_id={self.reference_id}, metadata={self.metadata}, client={self._client})"
 
     def __eq__(self, other):
         return self.id == other.id

tests/__init__.py

@@ -17,6 +17,7 @@
     DATASET_ID_KEY,
 )
 
+
 @pytest.fixture()
 def dataset(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
@@ -25,6 +26,7 @@ def dataset(CLIENT):
     response = CLIENT.delete_dataset(ds.id)
     assert response == {}
 
+
 def test_dataset_create_and_delete(CLIENT):
     # Creation
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
@@ -88,14 +90,17 @@ def test_dataset_list_autotags(CLIENT, dataset):
     autotag_response = CLIENT.list_autotags(dataset.id)
     assert autotag_response == []
 
+
 def test_slice_create_and_delete(dataset):
     # Dataset upload
     ds_items = []
     for url in TEST_IMG_URLS:
-        ds_items.append(DatasetItem(
-            image_location=url,
-            reference_id=reference_id_from_url(url),
-        ))
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
     response = dataset.append(ds_items)
     assert ERROR_PAYLOAD not in response.json()
 
@@ -121,10 +126,12 @@ def test_slice_append(dataset):
     # Dataset upload
     ds_items = []
     for url in TEST_IMG_URLS:
-        ds_items.append(DatasetItem(
-            image_location=url,
-            reference_id=reference_id_from_url(url),
-        ))
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
     response = dataset.append(ds_items)
     assert ERROR_PAYLOAD not in response.json()
 

tests/test_dataset.py

@@ -16,7 +16,7 @@
     ERROR_PAYLOAD,
     DATASET_ID_KEY,
 )
-from .helpers import (
+from helpers import (
     TEST_MODEL_NAME,
     TEST_MODEL_REFERENCE,
     TEST_MODEL_RUN,