Upload app.py

app.py

@@ -0,0 +1,117 @@
+import gradio as gr
+from transformers import SamModel, SamProcessor
+from PIL import Image, ImageDraw
+import torch
+import numpy as np
+
+# Load SAM
+processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+model = SamModel.from_pretrained("facebook/sam-vit-huge")
+model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Global variables
+global_state = {
+    "image": None,
+    "clicks": [],  # List of tuples: (x, y, label) where label=1 fg, 0 bg
+    "bbox": None,  # (x0, y0, x1, y1)
+}
+
+# Helper to apply mask overlay
+def apply_mask_overlay(image: Image.Image, mask: np.ndarray, color=(255, 0, 0)) -> Image.Image:
+    mask_img = Image.fromarray(mask.astype(np.uint8) * 255).convert("L")
+    color_mask = Image.new("RGB", image.size, color)
+    mask_rgb = Image.composite(color_mask, image, mask_img)
+    blended = Image.blend(image, mask_rgb, alpha=0.5)
+    return blended
+
+# Set image
+def upload_image(img):
+    global_state["image"] = img
+    global_state["clicks"] = []
+    global_state["bbox"] = None
+    return "Image uploaded. Now click or draw a box.", img
+
+# Handle point clicks
+def on_click(evt: gr.SelectData):
+    if global_state["image"] is None:
+        return gr.update(), "Please upload an image first."
+
+    # Default to foreground click
+    global_state["clicks"].append((evt.index[0], evt.index[1], 1))
+    return gr.update(), f"Point added: ({evt.index[0]}, {evt.index[1]})"
+
+# Handle bounding box input
+def set_bbox(x0: int, y0: int, x1: int, y1: int):
+    global_state["bbox"] = (x0, y0, x1, y1)
+    return f"Bounding box set: ({x0}, {y0}, {x1}, {y1})"
+
+# Run segmentation
+def run_segmentation():
+    if global_state["image"] is None:
+        return None, "Please upload an image first."
+
+    image = global_state["image"]
+    inputs = processor(image, return_tensors="pt").to(device)
+
+    if global_state["clicks"]:
+        points = torch.tensor([[[x, y] for (x, y, l) in global_state["clicks"]]], device=device)
+        labels = torch.tensor([[l for (_, _, l) in global_state["clicks"]]], device=device)
+        inputs.update({"input_points": points, "input_labels": labels})
+
+    if global_state["bbox"]:
+        x0, y0, x1, y1 = global_state["bbox"]
+        box = torch.tensor([[[x0, y0, x1, y1]]], device=device)
+        inputs.update({"input_boxes": box})
+
+    with torch.no_grad():
+        outputs = model(**inputs, multimask_output=False)
+
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )[0]
+
+    final_mask = masks[0].numpy()
+    overlayed = apply_mask_overlay(image.convert("RGB"), final_mask)
+
+    return overlayed, "Segmentation complete."
+
+# Reset
+def reset_all():
+    global_state["image"] = None
+    global_state["clicks"] = []
+    global_state["bbox"] = None
+    return None, None, "State reset."
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🔬 Interactive Pathology Segmentation with SAM")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Pathology Image")
+            upload_status = gr.Textbox(label="Status")
+            upload_btn = gr.Button("Upload & Set Image")
+            click_output = gr.Textbox(label="Click Info")
+            run_btn = gr.Button("Run Segmentation")
+            reset_btn = gr.Button("Reset")
+            bbox_coords = gr.Textbox(label="Manual BBox (x0,y0,x1,y1)")
+            set_bbox_btn = gr.Button("Set Bounding Box")
+        with gr.Column():
+            image_output = gr.Image(type="pil", label="Segmentation Output")
+
+    # Handlers
+    upload_btn.click(upload_image, inputs=image_input, outputs=[upload_status, image_output])
+    image_input.select(on_click, outputs=[image_output, click_output])
+    run_btn.click(run_segmentation, outputs=[image_output, upload_status])
+    reset_btn.click(reset_all, outputs=[image_output, click_output, upload_status])
+    set_bbox_btn.click(
+        lambda coords: set_bbox(*map(int, coords.split(","))),
+        inputs=bbox_coords,
+        outputs=click_output
+    )
+
+demo.launch()