Upload 5 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+003_img.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,96 @@
+import gradio as gr
+from gradio_bbox_annotator import BBoxAnnotator
+from PIL import Image
+import numpy as np
+
+# 你已有的推理代码
+from inference import load_model, get_embedding, run
+
+# ---- 仅加载一次模型 ----
+model, device = load_model("medsam_vit_b.pth")
+def predict(value):
+    # value: (image_path, [(xmin, ymin, xmax, ymax, label), ...])
+    return value  # 直接回显
+
+def make_example(path):
+    return [path, []]
+
+def parse_first_bbox(bboxes):
+    """
+    从 annot 的 bboxes 里取第一个框，返回 (xmin, ymin, xmax, ymax)
+    兼容两种格式：
+    - dict: {"x":..,"y":..,"width":..,"height":..}
+    - list: [xmin, ymin, xmax, ymax, ...]
+    """
+    if not bboxes:
+        return None
+    b = bboxes[0]
+    if isinstance(b, dict):
+        x, y = float(b["x"]), float(b["y"])
+        w, h = float(b["width"]), float(b["height"])
+        return x, y, x + w, y + h
+    if isinstance(b, (list, tuple)) and len(b) >= 4:
+        return float(b[0]), float(b[1]), float(b[2]), float(b[3])
+    return None
+
+def segment(annot_value):
+    """
+    annot_value 形如 [image_path, bboxes]
+    - image_path: 字符串
+    - bboxes: 框列表
+    """
+    if annot_value is None or len(annot_value) < 1:
+        return None, "请先在上方上传图片并拖一个矩形框。"
+
+    img_path = annot_value[0]
+    bboxes = annot_value[1] if len(annot_value) > 1 else []
+
+    if not bboxes:
+        return None, "未检测到矩形框，请在标注区按住左键拖拽一个框。"
+
+    # 读取图片
+    img = Image.open(img_path).convert("RGB")
+    img_np = np.array(img)
+    H, W, _ = img_np.shape
+
+    # 取第一个框
+    box = parse_first_bbox(bboxes)
+    if box is None:
+        return None, "解析矩形框失败，请重画。"
+
+    xmin, ymin, xmax, ymax = box
+
+    # 归一化到 1024 并推理
+    box_np = np.array([[xmin, ymin, xmax, ymax]], dtype=float)
+    box_1024 = box_np / np.array([W, H, W, H]) * 1024.0
+
+    embedding = get_embedding(model, img_np, device)
+    mask = run(model, embedding, box_1024, H, W)  # (H, W) 0/1
+
+    # 黑白 mask（白=前景）
+    mask_rgb = np.stack([mask * 255] * 3, axis=-1).astype(np.uint8)
+    bbox_text = f"xmin={int(xmin)}, ymin={int(ymin)}, xmax={int(xmax)}, ymax={int(ymax)}"
+
+    return Image.fromarray(mask_rgb), bbox_text
+
+# --- 构造一个可用的示例值（让画布里有图可直接拖） ---
+example = ("003_img.png", [(50, 60, 120, 150, "cell")])
+
+demo = gr.Interface(
+    fn=segment,  # ← 调你的推理函数
+    inputs=BBoxAnnotator(
+        value=example,                 # 默认示例；组件里自带“上传”按钮，可以换图
+        categories=["cell", "nucleus"],
+        label="upload"
+    ),
+    outputs=[
+        gr.Image(type="pil", label="Mask result"),
+        gr.Textbox(label="location")
+    ],
+    examples=[[example]],
+    cache_examples=False
+)
+
+if __name__ == "__main__":
+    demo.launch(server_port=7860)
+

inference.py

@@ -0,0 +1,48 @@
+import torch
+import numpy as np
+from skimage import transform
+from segment_anything import sam_model_registry
+
+MEDSAM_IMG_INPUT_SIZE = 1024
+
+def load_model(checkpoint_path):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = sam_model_registry["vit_b"](checkpoint=checkpoint_path).to(device)
+    model.eval()
+    return model, device
+
+@torch.no_grad()
+def get_embedding(model, img_np, device):
+    img_1024 = transform.resize(
+        img_np, (1024, 1024), order=3, preserve_range=True, anti_aliasing=True
+    ).astype(np.uint8)
+
+    img_1024 = (img_1024 - img_1024.min()) / np.clip(img_1024.max() - img_1024.min(), 1e-8, None)
+    img_tensor = torch.tensor(img_1024).float().permute(2, 0, 1).unsqueeze(0).to(device)
+    return model.image_encoder(img_tensor)
+
+@torch.no_grad()
+def run(model, embedding, box_1024, H, W):
+    box_torch = torch.as_tensor(box_1024, dtype=torch.float, device=embedding.device)
+    if len(box_torch.shape) == 2:
+        box_torch = box_torch[:, None, :]  # (B, 1, 4)
+
+    sparse_embeddings, dense_embeddings = model.prompt_encoder(
+        points=None,
+        boxes=box_torch,
+        masks=None,
+    )
+    low_res_logits, _ = model.mask_decoder(
+        image_embeddings=embedding,
+        image_pe=model.prompt_encoder.get_dense_pe(),
+        sparse_prompt_embeddings=sparse_embeddings,
+        dense_prompt_embeddings=dense_embeddings,
+        multimask_output=False,
+    )
+
+    low_res_pred = torch.sigmoid(low_res_logits)
+    low_res_pred = torch.nn.functional.interpolate(
+        low_res_pred, size=(H, W), mode="bilinear", align_corners=False
+    )
+    low_res_pred = low_res_pred.squeeze().cpu().numpy()
+    return (low_res_pred > 0.5).astype(np.uint8)

medsam_vit_b.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34b34b78c1d18cb8c6bf84cf9c00e135d6d6c965699f3c0e31ef1bc9dcb5be74
+size 375049145

requirements.txt

@@ -0,0 +1,6 @@
+streamlit
+torch
+numpy
+Pillow
+scikit-image
+