Upload app.py

app.py

@@ -8,6 +8,8 @@ from typing import Tuple, Dict
 from timeit import default_timer as timer
 from skimage import io, transform
 import os
+import base64
+import json
 
 import torch.nn.functional as F
 
@@ -104,7 +106,7 @@ def predict(img) -> Tuple[Dict, float]:
     with torch.inference_mode():
       image_embedding = medsam_model.image_encoder(img_1024_tensor)  # (1, 256, 64, 64)
     # define the inputbox  
-    input_box = np.array([[425, 200, 600, 400]])
+    input_box = np.array([[125, 275, 190, 350]])
     # transfer box_np t0 1024x1024 scale
     box_1024 = input_box / np.array([W, H, W, H]) * 1024
 
@@ -120,10 +122,13 @@ def predict(img) -> Tuple[Dict, float]:
     show_box(input_box[0], ax[1])
     ax[1].set_title("MedSAM Segmentation")
     # Calculate the prediction time
-    
+    image_embedding = image_embedding.cpu().numpy().tobytes()
+
+    # Serialize the response data to JSON format
+    serialized_data = json.dumps([base64.b64encode(image_embedding).decode('ascii')])
 
     # Return the prediction dictionary and prediction time
-    return fig, pred_time
+    return fig, pred_time,serialized_data
 
 # 4. Gradio app
 # Create title, description and article strings
@@ -138,7 +143,8 @@ example_list = [["examples/" + example] for example in os.listdir("examples")]
 demo = gr.Interface(fn=predict, # mapping function from input to output
                     inputs=gr.Image(type="pil"), # what are the inputs?
                     outputs=[gr.Plot(label="Predictions"), # what are the outputs?
-                            gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
+                            gr.Number(label="Prediction time (s)"),
+                            gr.JSON(label="Embedding Image")], # our fn has two outputs, therefore we have two outputs
                     examples=example_list,
                     title=title,
                     description=description,