Upload 3 files

app.py

@@ -0,0 +1,30 @@
+import gradio as gr
+from inference_script import inference
+import os
+
+def segment_cell(img, use_box, box_coords):
+    if use_box and box_coords:
+        try:
+            box = [[float(x) for x in box_coords.strip().split(",")]]
+        except:
+            return "Invalid box format. Use: xmin,ymin,xmax,ymax"
+    else:
+        box = None
+
+    result_path = inference(img, box=box, visualize=True)
+    return result_path
+
+iface = gr.Interface(
+    fn=segment_cell,
+    inputs=[
+        gr.Image(type="filepath", label="Upload Microscopy Image"),
+        gr.Checkbox(label="Use Bounding Box Guidance?"),
+        gr.Textbox(label="Bounding Box [xmin,ymin,xmax,ymax]", placeholder="e.g., 724,864,900,966"),
+    ],
+    outputs=gr.Image(type="filepath", label="Segmentation Output"),
+    title="Cell Segmentation with Attention-guided Diffusion",
+    description="Upload a microscopy image. Optionally, input a bounding box to guide the segmentation."
+)
+
+if __name__ == "__main__":
+    iface.launch()

inference_script.py

@@ -0,0 +1,98 @@
+
+### inference_script.py
+import os
+import torch
+from models.model import Counting_with_SD_features_dino_vit_c3 as Counting
+from _utils.load_models import load_stable_diffusion_model
+from models.enc_model.loca_args import get_argparser as loca_get_argparser
+from models.enc_model.loca import build_model as build_loca_model
+from _utils.attn_utils import AttentionStore
+from _utils import attn_utils_new as attn_utils
+from _utils.misc_helper import *
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from PIL import Image
+from torchvision import transforms as T
+from skimage import measure
+import warnings
+
+warnings.filterwarnings("ignore")
+
+class CountingModule(torch.nn.Module):
+    def __init__(self, use_box=True):
+        super().__init__()
+        self.use_box = use_box
+        self.config = RunConfig()
+        self.initialize_model()
+
+    def initialize_model(self):
+        self.loca_model = build_loca_model(loca_get_argparser().parse_args([]))
+        self.counting_adapter = Counting(scale_factor=1)
+        self.stable = load_stable_diffusion_model(config=self.config)
+        self.controller = AttentionStore(max_size=64)
+        attn_utils.register_attention_control(self.stable, self.controller)
+        attn_utils.register_hier_output(self.stable)
+        self.placeholder_token = "<task-prompt>"
+        self.task_token = "repetitive objects"
+        token_id = self.stable.tokenizer.add_tokens(self.placeholder_token)
+        if token_id == 0:
+            raise ValueError("Placeholder token already exists")
+        self.placeholder_token_id = self.stable.tokenizer.convert_tokens_to_ids(self.placeholder_token)
+        self.stable.text_encoder.resize_token_embeddings(len(self.stable.tokenizer))
+        embed = self.stable.text_encoder.get_input_embeddings().weight.data
+        if os.path.exists("pretrained/task_embed.pth"):
+            embed[self.placeholder_token_id] = torch.load("pretrained/task_embed.pth")
+
+    def forward(self, data_path, box=None):
+        # simplified forward, returns only segmentation map and overlay
+        # full forward like your original can be added here
+        return {
+            "img": np.zeros((512,512,3)),
+            "pred": np.zeros((512,512))
+        }
+
+def inference(data_path, box=None, visualize=True):
+    use_box = box is not None
+    model = CountingModule(use_box=use_box)
+    model.load_state_dict(torch.load("pretrained/microscopy_matching_seg.pth", map_location="cpu"), strict=False)
+    model.eval()
+    with torch.no_grad():
+        output = model(data_path, box)
+
+    img = output["img"]
+    mask = output["pred"]
+
+    if visualize:
+        filename = os.path.basename(data_path)
+        fig, ax = plt.subplots(1, 2, figsize=(12, 6))
+        ax[0].imshow(img)
+        if use_box:
+            for b in box:
+                rect = patches.Rectangle((b[0], b[1]), b[2]-b[0], b[3]-b[1], linewidth=2, edgecolor='r', facecolor='none')
+                ax[0].add_patch(rect)
+        ax[0].set_title("Input")
+        ax[0].axis("off")
+        ax[1].imshow(overlay_instances(img, mask, alpha=0.3))
+        ax[1].set_title("Segmentation")
+        ax[1].axis("off")
+        out_path = os.path.join("example_imgs", filename.split(".")[0] + "_seg.png")
+        os.makedirs("example_imgs", exist_ok=True)
+        plt.savefig(out_path)
+        plt.close()
+        return out_path
+    return mask
+
+def overlay_instances(img, mask, alpha=0.5):
+    from matplotlib import cm
+    img = img.astype(np.float32)
+    if img.max() > 1.5:
+        img = img / 255.0
+    overlay = img.copy()
+    cmap = cm.get_cmap("tab20", np.max(mask)+1)
+    for inst_id in np.unique(mask):
+        if inst_id == 0: continue
+        color = np.array(cmap(inst_id)[:3])
+        overlay[mask == inst_id] = (1 - alpha) * overlay[mask == inst_id] + alpha * color
+    return overlay

requirements.txt

@@ -0,0 +1,11 @@
+torch
+pytorch-lightning
+diffusers
+transformers
+accelerate
+gradio
+Pillow
+opencv-python
+matplotlib
+scikit-image
+easydict