Update app.py

app.py

@@ -50,15 +50,20 @@ def predict_segmentation(image):
     # Create an alpha channel: 255 where there is segmentation, 0 otherwise
     rgba_img[:, :, 3] = np.where(predicted_img_resized > 0, 255, 0)
 
-    # Convert the numpy array to an image and return it
+    # Convert the numpy array to an image
     output_image = Image.fromarray(rgba_img)
-    return output_image
+
+    # Save the image as PNG to return it
+    output_image_path = "/tmp/segmented_output.png"
+    output_image.save(output_image_path)
+
+    return output_image_path
 
 # Gradio Interface
 iface = gr.Interface(
     fn=predict_segmentation,
     inputs="image",
-    outputs="image",
+    outputs="file",  # Return the file path to download the PNG
     live=False
 )