Dict return

app.py

@@ -11,12 +11,15 @@ from models.preprocess import preprocess
 
 FAST_SAM = loadModel()
 
+
 def base64_to_image(base64_str):
     image_data = base64.b64decode(base64_str)
     image = Image.open(BytesIO(image_data))
     return image
 
 # Main processing function
+
+
 def segment_marker(img_rgb: Image.Image, marker_coordinates: str):
     # Parse marker coordinates from JSON string
     try:
@@ -39,11 +42,10 @@ def segment_marker(img_rgb: Image.Image, marker_coordinates: str):
         bg_only_removed_img = removeOnlyBg(img_rgb, masks[0])
         img_base64_only_bg = convertToBuffer(bg_only_removed_img)
 
-        # Convert base64 strings to PIL images for Gradio
-        img_bg_segmented = base64_to_image(img_base64_bg_segmented)
-        img_bg_only_removed = base64_to_image(img_base64_only_bg)
-
-        return img_bg_segmented, img_bg_only_removed  # Return as two separate images
+        return {
+            'bg_removed_segmented_img': f'data:image/png;base64,{img_base64_bg_segmented}',
+            'bg_only_removed_segmented_img': f'data:image/png;base64,{img_base64_only_bg}'
+        }
 
     except Exception as e:
         print(f"An error occurred: {str(e)}")
@@ -57,13 +59,9 @@ iface = gr.Interface(
         gr.Image(type="pil", label="Upload Image"),
         gr.Textbox(label="Markers Coordinates (JSON format)")
     ],
-    outputs=[
-        gr.Image(type="pil", label="Background Removed with Segmentation"),
-        gr.Image(type="pil", label="Only Background Removed")
-    ],
+    outputs="json",  # Set output to JSON format to return the dictionary
     title="Image Segmentation with Background Removal",
     description="Upload an image and JSON-formatted marker coordinates to perform image segmentation and background removal."
 )
-
 # Run the Gradio app
 iface.launch(share=True)