Update app.py

app.py

@@ -19,11 +19,19 @@ def segment_image(input_image, text_prompt):
     # Get the predicted segmentation
     preds = outputs.logits.squeeze().sigmoid()
 
-    # Convert the prediction to a PIL image
-    segmentation = (preds > 0.5).float()
-    segmentation_image = Image.fromarray((segmentation.numpy() * 255).astype(np.uint8))
+    # Convert the prediction to a numpy array and scale to 0-255
+    segmentation = (preds.numpy() * 255).astype(np.uint8)
 
-    return segmentation_image
+    # Create a colored heatmap
+    heatmap = np.zeros((segmentation.shape[0], segmentation.shape[1], 3), dtype=np.uint8)
+    heatmap[:, :, 0] = segmentation  # Red channel
+    heatmap[:, :, 2] = 255 - segmentation  # Blue channel
+
+    # Blend the heatmap with the original image
+    original_image = np.array(input_image)
+    blended = (0.7 * original_image + 0.3 * heatmap).astype(np.uint8)
+
+    return Image.fromarray(blended)
 
 # Create Gradio interface
 iface = gr.Interface(