change payload to always take a list of points of 1 or more for all p…

…rompts, fix missing args to decode_Single_point
developmentseed · Nov 30, 2023 · 73d9b97 · 73d9b97
1 parent 16e79a4
commit 73d9b97
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 110 deletions.
diff --git a/handler_decode.py b/handler_decode.py
@@ -6,7 +6,7 @@
     prepare_decode_inputs,
     decode_single_point,
     decode_multi_point,
-    decode_multi_split,
+    decode_multi_point_split,
     mask_to_geojson,
     masks_to_utf8,
 )
@@ -33,11 +33,13 @@ def preprocess(self, data):
 
     def inference(self, payload):
         if payload["decode_type"] == "single_point":
-            return decode_single_point(payload, self.ort_session)
+            input_label = payload['input_label'][0]
+            input_point = payload['input_prompt'][0]
+            return decode_single_point(payload, self.ort_session, input_point, input_label)
         elif payload["decode_type"] == "multi_point":
             return decode_multi_point(payload, self.ort_session)
         elif payload["decode_type"] == "bbox":
-            return decode_multi_split(payload, self.ort_session)
+            return decode_multi_point_split(payload, self.ort_session)
 
     def handle(self, data, context):
         """

diff --git a/notebooks/test_sam_endpoints.ipynb b/notebooks/test_sam_endpoints.ipynb
diff --git a/notebooks/test_sam_endpoints_multi.ipynb b/notebooks/test_sam_endpoints_multi.ipynb
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -69,8 +69,8 @@ def single_point_payload(image_embeddings_as_str):
             "data": {
                 "image_embeddings": image_embeddings_as_str,
                 "image_shape": [512, 512],
-                "input_prompt": input_point_on_slick,
-                "input_label": 1,
+                "input_prompt": [input_point_on_slick],
+                "input_label": [1],
                 "decode_type": "single_point",
             }
         }