Merge pull request #73 from scaleapi/classpdf

Add prediction class pdf support

Author: jihan-yin

nucleus/constants.py

@@ -11,6 +11,7 @@
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
 
+CLASS_PDF_KEY = "class_pdf"
 CONFIDENCE_KEY = "confidence"
 DATASET_ID_KEY = "dataset_id"
 DATASET_ITEM_IDS_KEY = "dataset_item_ids"

nucleus/prediction.py

@@ -16,6 +16,7 @@
     Y_KEY,
     WIDTH_KEY,
     HEIGHT_KEY,
+    CLASS_PDF_KEY,
     CONFIDENCE_KEY,
     VERTICES_KEY,
     ANNOTATIONS_KEY,
@@ -54,6 +55,7 @@ def __init__(
         confidence: Optional[float] = None,
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
     ):
         super().__init__(
             label=label,
@@ -67,11 +69,13 @@ def __init__(
             metadata=metadata,
         )
         self.confidence = confidence
+        self.class_pdf = class_pdf
 
     def to_payload(self) -> dict:
         payload = super().to_payload()
         if self.confidence is not None:
             payload[CONFIDENCE_KEY] = self.confidence
+            payload[CLASS_PDF_KEY] = self.class_pdf
 
         return payload
 
@@ -89,6 +93,7 @@ def from_json(cls, payload: dict):
             confidence=payload.get(CONFIDENCE_KEY, None),
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
+            class_pdf=payload.get(CLASS_PDF_KEY, {}),
         )
 
 
@@ -102,6 +107,7 @@ def __init__(
         confidence: Optional[float] = None,
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
     ):
         super().__init__(
             label=label,
@@ -112,11 +118,13 @@ def __init__(
             metadata=metadata,
         )
         self.confidence = confidence
+        self.class_pdf = class_pdf
 
     def to_payload(self) -> dict:
         payload = super().to_payload()
         if self.confidence is not None:
             payload[CONFIDENCE_KEY] = self.confidence
+            payload[CLASS_PDF_KEY] = self.class_pdf
 
         return payload
 
@@ -131,4 +139,5 @@ def from_json(cls, payload: dict):
             confidence=payload.get(CONFIDENCE_KEY, None),
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
+            class_pdf=payload.get(CLASS_PDF_KEY, {}),
         )

tests/helpers.py

@@ -87,13 +87,31 @@ def reference_id_from_url(url):
 ]
 TEST_SEGMENTATION_PREDICTIONS = TEST_SEGMENTATION_ANNOTATIONS
 
+TEST_BOX_MODEL_PDF = {
+    box_annotation["label"]: 1 / len(TEST_BOX_ANNOTATIONS)
+    for box_annotation in TEST_BOX_ANNOTATIONS
+}
+
+TEST_POLYGON_MODEL_PDF = {
+    polygon_annotation["label"]: 1 / len(TEST_POLYGON_ANNOTATIONS)
+    for polygon_annotation in TEST_POLYGON_ANNOTATIONS
+}
+
 TEST_BOX_PREDICTIONS = [
-    {**TEST_BOX_ANNOTATIONS[i], "confidence": 0.10 * i}
+    {
+        **TEST_BOX_ANNOTATIONS[i],
+        "confidence": 0.10 * i,
+        "class_pdf": TEST_BOX_MODEL_PDF,
+    }
     for i in range(len(TEST_BOX_ANNOTATIONS))
 ]
 
 TEST_POLYGON_PREDICTIONS = [
-    {**TEST_POLYGON_ANNOTATIONS[i], "confidence": 0.10 * i}
+    {
+        **TEST_POLYGON_ANNOTATIONS[i],
+        "confidence": 0.10 * i,
+        "class_pdf": TEST_POLYGON_MODEL_PDF,
+    }
     for i in range(len(TEST_POLYGON_ANNOTATIONS))
 ]
 

tests/test_dataset.py

@@ -275,7 +275,6 @@ def test_annotate_async(dataset: Dataset):
                 "processed": 2,
             },
             "segmentation_upload": {
-                "errors": [],
                 "ignored": 0,
                 "n_errors": 0,
                 "processed": 1,
@@ -311,7 +310,6 @@ def test_annotate_async_with_error(dataset: Dataset):
                 "processed": 1,
             },
             "segmentation_upload": {
-                "errors": [],
                 "ignored": 0,
                 "n_errors": 0,
                 "processed": 1,

tests/test_prediction.py

@@ -297,7 +297,6 @@ def test_mixed_pred_upload_async(model_run: ModelRun):
                 "processed": 2,
             },
             "segmentation_upload": {
-                "errors": [],
                 "ignored": 0,
                 "n_errors": 0,
                 "processed": 1,
@@ -334,7 +333,6 @@ def test_mixed_pred_upload_async_with_error(model_run: ModelRun):
                 "processed": 1,
             },
             "segmentation_upload": {
-                "errors": [],
                 "ignored": 0,
                 "n_errors": 0,
                 "processed": 1,