pytests working

Author: jihan-yin

nucleus/__init__.py

@@ -572,28 +572,24 @@ def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
     def predict(
         self,
         model_run_id: str,
-        payload: Dict[str, List[Union[BoxPrediction, PolygonPrediction]]],
+        annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        update: bool,
         batch_size: int = 100,
     ):
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
-        :param payload:
-        {
-            "annotations": List[Union[Box2DPrediction, Polygon2DPrediction]],
-        }
+        :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        :param update: bool
         :return:
         {
             "dataset_id": str,
             "model_run_id": str,
             "annotations_processed: int,
         }
         """
-        predictions: List[Union[BoxPrediction, PolygonPrediction]] = payload[
-            ANNOTATIONS_KEY
-        ]
         batches = [
-            predictions[i : i + batch_size]
-            for i in range(0, len(predictions), batch_size)
+            annotations[i : i + batch_size]
+            for i in range(0, len(annotations), batch_size)
         ]
 
         agg_response = {
@@ -604,7 +600,11 @@ def predict(
         tqdm_batches = self.tqdm_bar(batches)
 
         for batch in tqdm_batches:
-            batch_payload = {ANNOTATIONS_KEY: batch}
+            batch_payload = construct_box_predictions_payload(
+                annotations,
+                update,
+            )
+            print(batch_payload)
             response = self._make_request(
                 batch_payload, f"modelRun/{model_run_id}/predict"
             )
@@ -701,7 +701,7 @@ def predictions_ref_id(self, model_run_id: str, ref_id: str):
         :param reference_id: reference_id of a dataset item.
         :return:
         {
-            "annotations": List[Box2DPrediction],
+            "annotations": List[BoxPrediction],
         }
         """
         return self._make_request(
@@ -726,7 +726,7 @@ def predictions_iloc(self, model_run_id: str, i: int):
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
         :return:
         {
-            "annotations": List[Box2DPrediction],
+            "annotations": List[BoxPrediction],
         }
         """
         return self._make_request(
@@ -755,7 +755,7 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
         :param dataset_item_id: dataset_item_id of a dataset item.
         :return:
         {
-            "annotations": List[Box2DPrediction],
+            "annotations": List[BoxPrediction],
         }
         """
         return self._make_request(

nucleus/model_run.py

@@ -1,5 +1,5 @@
 from typing import Optional, List, Dict, Any, Union
-from .constants import ANNOTATIONS_KEY
+from .constants import ANNOTATIONS_KEY, DEFAULT_ANNOTATION_UPDATE_MODE
 from .prediction import BoxPrediction, PolygonPrediction
 from .payload_constructor import construct_box_predictions_payload
 
@@ -61,30 +61,26 @@ def commit(self, payload: Optional[dict] = None) -> dict:
         return self._client.commit_model_run(self.model_run_id, payload)
 
     def predict(
-        self, annotations: List[Union[BoxPrediction, PolygonPrediction]]
+        self,
+        annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
     ) -> dict:
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
         :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
         :return:
         {
-            "dataset_id": str,
             "model_run_id": str,
-            "annotations_processed: int,
+            "predictions_processed: int,
         }
         """
-        payload: Dict[str, List[Any]] = construct_box_predictions_payload(
-            annotations
-        )
-        return self._client.predict(self.model_run_id, payload)
+        return self._client.predict(self.model_run_id, annotations, update)
 
     def iloc(self, i: int):
         """
         Returns Model Run Info For Dataset Item by its number.
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
-        :return:
-        {
-            "annotations": List[Union[BoxPrediction, PolygonPrediction]],
+        :return: List[Union[BoxPrediction, PolygonPrediction]],
         }
         """
         response = self._client.predictions_iloc(self.model_run_id, i)
@@ -94,10 +90,7 @@ def refloc(self, reference_id: str):
         """
         Returns Model Run Info For Dataset Item by its reference_id.
         :param reference_id: reference_id of a dataset item.
-        :return:
-        {
-            "annotations": List[Union[BoxPrediction, PolygonPrediction]],
-        }
+        :return: List[Union[BoxPrediction, PolygonPrediction]],
         """
         response = self._client.predictions_ref_id(
             self.model_run_id, reference_id

nucleus/payload_constructor.py

@@ -41,12 +41,13 @@ def construct_annotation_payload(
 
 def construct_box_predictions_payload(
     box_predictions: List[Union[BoxPrediction, PolygonPrediction]],
+    update: bool,
 ) -> dict:
     predictions = []
     for prediction in box_predictions:
         predictions.append(prediction.to_payload())
 
-    return {ANNOTATIONS_KEY: predictions}
+    return {ANNOTATIONS_KEY: predictions, ANNOTATION_UPDATE_KEY: update}
 
 
 def construct_model_creation_payload(

nucleus/prediction.py

@@ -37,7 +37,7 @@ def __init__(
 
     def to_payload(self) -> dict:
         payload = super().to_payload()
-        if self.confidence:
+        if self.confidence is not None:
             payload[CONFIDENCE_KEY] = self.confidence
 
         return payload
@@ -78,7 +78,7 @@ def __init__(
 
     def to_payload(self) -> dict:
         payload = super().to_payload()
-        if self.confidence:
+        if self.confidence is not None:
             payload[CONFIDENCE_KEY] = self.confidence
 
         return payload

tests/test_annotation.py

@@ -54,7 +54,6 @@ def test_polygon_gt_upload(dataset):
     response = dataset.refloc(annotation.reference_id)['annotations']
     assert len(response) == 1
     response_annotation = response[0]
-    print(response_annotation)
     assert_polygon_annotation_matches_dict(response_annotation, TEST_POLYGON_ANNOTATIONS[0])
 
 

tests/test_dataset.py

@@ -114,7 +114,6 @@ def test_slice_create_and_delete(dataset):
     assert len(response["reference_ids"]) == 2
     for item in ds_items[:2]:
         assert item.reference_id in response["reference_ids"]
-    print(response)
 
 
 def test_slice_append(dataset):

tests/test_prediction.py

@@ -4,6 +4,7 @@
     TEST_DATASET_NAME,
     TEST_MODEL_NAME,
     TEST_MODEL_REFERENCE,
+    TEST_MODEL_RUN,
     TEST_IMG_URLS,
     TEST_BOX_PREDICTIONS,
     TEST_POLYGON_PREDICTIONS,
@@ -16,7 +17,7 @@
 from nucleus.constants import ERROR_PAYLOAD
 
 @pytest.fixture()
-def dataset(CLIENT):
+def model_run(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
     ds_items = []
     for url in TEST_IMG_URLS:
@@ -30,46 +31,48 @@ def dataset(CLIENT):
 
     model = CLIENT.add_model(
         name=TEST_MODEL_NAME,
-        reference_id=TEST_MODEL_NAME
+        reference_id=TEST_MODEL_REFERENCE
     )
 
+    run = model.create_run(
+        name=TEST_MODEL_RUN,
+        dataset=ds,
+        predictions=[])
 
-    yield ds
+    yield run
 
     response = CLIENT.delete_dataset(ds.id)
     assert response == {}
+    response = CLIENT.delete_model(model.id)
+    assert response == {}
 
-
-def test_box_pred_upload(dataset):
+def test_box_pred_upload(model_run):
     prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
-    response = dataset.annotate(predictions=[prediction])
+    response = model_run.predict(annotations=[prediction])
 
-    assert response['dataset_id'] == dataset.id
+    assert response['model_run_id'] == model_run.model_run_id
     assert response['predictions_processed'] == 1
 
-    response = dataset.refloc(prediction.reference_id)['predictions']
+    response = model_run.refloc(prediction.reference_id)
     assert len(response) == 1
-    response_prediction = response[0]
-    assert_box_prediction_matches_dict(response_prediction, TEST_BOX_PREDICTIONS[0])
+    assert_box_prediction_matches_dict(response[0], TEST_BOX_PREDICTIONS[0])
 
 
-def test_polygon_pred_upload(dataset):
+def test_polygon_pred_upload(model_run):
     prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
-    response = dataset.annotate(predictions=[prediction])
+    response = model_run.predict(annotations=[prediction])
 
-    assert response['dataset_id'] == dataset.id
+    assert response['model_run_id'] == model_run.model_run_id
     assert response['predictions_processed'] == 1
 
-    response = dataset.refloc(prediction.reference_id)['predictions']
+    response = model_run.refloc(prediction.reference_id)
     assert len(response) == 1
-    response_prediction = response[0]
-    print(response_prediction)
-    assert_polygon_prediction_matches_dict(response_prediction, TEST_POLYGON_PREDICTIONS[0])
+    assert_polygon_prediction_matches_dict(response[0], TEST_POLYGON_PREDICTIONS[0])
 
 
-def test_box_pred_upload_update(dataset):
+def test_box_pred_upload_update(model_run):
     prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
-    response = dataset.annotate(predictions=[prediction])
+    response = model_run.predict(annotations=[prediction])
 
     assert response['predictions_processed'] == 1
 
@@ -79,19 +82,18 @@ def test_box_pred_upload_update(dataset):
     prediction_update_params['reference_id'] = TEST_BOX_PREDICTIONS[0]['reference_id']
 
     prediction_update = BoxPrediction(**prediction_update_params)
-    response = dataset.annotate(predictions=[prediction_update], update=True)
+    response = model_run.predict(annotations=[prediction_update], update=True)
 
     assert response['predictions_processed'] == 1
 
-    response = dataset.refloc(prediction.reference_id)['predictions']
+    response = model_run.refloc(prediction.reference_id)
     assert len(response) == 1
-    response_prediction = response[0]
-    assert_box_prediction_matches_dict(response_prediction, prediction_update_params)
+    assert_box_prediction_matches_dict(response[0], prediction_update_params)
 
 
-def test_box_pred_upload_ignore(dataset):
+def test_box_pred_upload_ignore(model_run):
     prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
-    response = dataset.annotate(predictions=[prediction])
+    response = model_run.predict(annotations=[prediction])
 
     assert response['predictions_processed'] == 1
 
@@ -101,19 +103,18 @@ def test_box_pred_upload_ignore(dataset):
     prediction_update_params['reference_id'] = TEST_BOX_PREDICTIONS[0]['reference_id']
     prediction_update = BoxPrediction(**prediction_update_params)
     # Default behavior is ignore.
-    response = dataset.annotate(predictions=[prediction_update])
+    response = model_run.predict(annotations=[prediction_update])
 
     assert response['predictions_processed'] == 1
 
-    response = dataset.refloc(prediction.reference_id)['predictions']
+    response = model_run.refloc(prediction.reference_id)
     assert len(response) == 1
-    response_prediction = response[0]
-    assert_box_prediction_matches_dict(response_prediction, TEST_BOX_PREDICTIONS[0])
+    assert_box_prediction_matches_dict(response[0], TEST_BOX_PREDICTIONS[0])
 
 
-def test_polygon_pred_upload_update(dataset):
+def test_polygon_pred_upload_update(model_run):
     prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
-    response = dataset.annotate(predictions=[prediction])
+    response = model_run.predict(annotations=[prediction])
 
     assert response['predictions_processed'] == 1
 
@@ -123,19 +124,18 @@ def test_polygon_pred_upload_update(dataset):
     prediction_update_params['reference_id'] = TEST_POLYGON_PREDICTIONS[0]['reference_id']
 
     prediction_update = PolygonPrediction(**prediction_update_params)
-    response = dataset.annotate(predictions=[prediction_update], update=True)
+    response = model_run.predict(annotations=[prediction_update], update=True)
 
     assert response['predictions_processed'] == 1
 
-    response = dataset.refloc(prediction.reference_id)['predictions']
+    response = model_run.refloc(prediction.reference_id)
     assert len(response) == 1
-    response_prediction = response[0]
-    assert_polygon_prediction_matches_dict(response_prediction, prediction_update_params)
+    assert_polygon_prediction_matches_dict(response[0], prediction_update_params)
 
 
-def test_polygon_pred_upload_ignore(dataset):
+def test_polygon_pred_upload_ignore(model_run):
     prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
-    response = dataset.annotate(predictions=[prediction])
+    response = model_run.predict(annotations=[prediction])
 
     assert response['predictions_processed'] == 1
 
@@ -146,11 +146,10 @@ def test_polygon_pred_upload_ignore(dataset):
 
     prediction_update = PolygonPrediction(**prediction_update_params)
     # Default behavior is ignore.
-    response = dataset.annotate(predictions=[prediction_update])
+    response = model_run.predict(annotations=[prediction_update])
 
     assert response['predictions_processed'] == 1
 
-    response = dataset.refloc(prediction.reference_id)['predictions']
+    response = model_run.refloc(prediction.reference_id)
     assert len(response) == 1
-    response_prediction = response[0]
-    assert_polygon_prediction_matches_dict(response_prediction, TEST_POLYGON_PREDICTIONS[0])
+    assert_polygon_prediction_matches_dict(response[0], TEST_POLYGON_PREDICTIONS[0])