push app

app.py

@@ -73,7 +73,7 @@ def load_models():
     return fastsam_model, sd_pipe
 
 # Initialize the models (this will run only once thanks to caching)
-# fastsam_model, sd_pipe = load_models()
+fastsam_model, sd_pipe = load_models()
 
 # Ensure we have a state for removing_dots
 if "is_removing_dot" not in st.session_state:
@@ -111,13 +111,38 @@ else:
     if "coords_list" not in st.session_state:
         st.session_state.coords_list = []
 
-    # Create a copy of the image and draw red dots for every stored coordinate.
-    img_with_dots = img.copy()
-    draw = ImageDraw.Draw(img_with_dots)
-    for coord in st.session_state.coords_list:
-        cx, cy = int(coord["x"]), int(coord["y"])
+    # --- Compute Segmentation Overlay ---
+    # If any points have been stored, run segmentation with FastSAM.
+    if st.session_state.coords_list:
+        points = [[int(pt["x"]), int(pt["y"])] for pt in st.session_state.coords_list]
+        labels = [1] * len(points)
+        results = fastsam_model(img, points=points, labels=labels)
+        # Assume results[0].masks.data is a tensor with shape (N, H, W)
+        masks_tensor = results[0].masks.data
+        masks = masks_tensor.cpu().numpy()
+        if masks.ndim == 3 and masks.shape[0] > 0:
+            # Combine masks (logical OR via max)
+            combined_mask = np.max(masks, axis=0)
+            combined_mask_img = Image.fromarray((combined_mask * 255).astype(np.uint8))
+            # Create a red overlay with transparency
+            overlay = Image.new("RGBA", img.size, (255, 0, 0, 100))
+            base = img.convert("RGBA")
+            mask_for_overlay = combined_mask_img.convert("L")
+            seg_overlay = Image.composite(overlay, base, mask_for_overlay)
+        else:
+            seg_overlay = img.copy()
+    else:
+        seg_overlay = img.copy()
+
+    # --- Draw Red Dots on Top ---
+    final_img = seg_overlay.copy()
+    draw = ImageDraw.Draw(final_img)
+    for pt in st.session_state.coords_list:
+        cx, cy = int(pt["x"]), int(pt["y"])
         draw.ellipse((cx - 5, cy - 5, cx + 5, cy + 5), fill="red")
 
+  
+
     # Use the interactive component as the display canvas, showing the image with all dots.
     new_coord = streamlit_image_coordinates(img_with_dots, key="click_img")
 
@@ -139,7 +164,7 @@ else:
     else:
         st.session_state.is_removing_dot = False
 
-    st.write("Coordinates:", st.session_state.get("coords"))
+    st.write("Stored coordinates:", st.session_state.coords_list)