Update app.py

app.py

@@ -3,37 +3,19 @@ from gradio_bbox_annotator import BBoxAnnotator
 from PIL import Image
 import numpy as np
 
-# 你已有的推理代码
 from inference import load_model, get_embedding, run
 import torch
 import os
 import spaces
 
-# @spaces.GPU
-# def warmup():
-#     import torch
-#     if torch.cuda.is_available():
-#         torch.zeros(1).to("cuda")
 
-
-# # ---- 仅加载一次模型 ----
-# 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, []]
-
-# --------- 全局状态 ---------
 MODEL = None
 DEVICE = torch.device("cpu")
 CUDA_READY = False
 
 def load_model_cpu(checkpoint_path: str):
     global MODEL, DEVICE
-    # 要求 inference.load_model 不要在内部 .to("cuda")
-    MODEL, _ = load_model(checkpoint_path)  # 或者直接返回 model
+    MODEL, _ = load_model(checkpoint_path)  
     MODEL = MODEL.to("cpu")
     MODEL.eval()
     DEVICE = torch.device("cpu")
@@ -44,22 +26,17 @@ load_model_cpu("medsam_vit_b.pth")
 def prepare_cuda():
     global MODEL, DEVICE, CUDA_READY
     if torch.cuda.is_available() and not CUDA_READY:
-        print("🎯 CUDA is available. Moving model to GPU...")
+        print("CUDA is available. Moving model to GPU...")
         MODEL.to("cuda")
         DEVICE = torch.device("cuda")
         CUDA_READY = True
         _ = torch.zeros(1, device=DEVICE)
-        print("✅ Model moved to CUDA.")
+        print("Model moved to CUDA.")
     else:
-        print("❌ CUDA not available or already initialized.")
+        print("CUDA not available or already initialized.")
 
 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]
@@ -73,52 +50,46 @@ def parse_first_bbox(bboxes):
 
 def segment(annot_value):
     prepare_cuda()
-    """
-    annot_value 形如 [image_path, bboxes]
-    - image_path: 字符串
-    - bboxes: 框列表
-    """
+
     if annot_value is None or len(annot_value) < 1:
-        return None, "请先在上方上传图片并拖一个矩形框。"
+        return None, 
 
     img_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
 
     if not bboxes:
-        return None, "未检测到矩形框，请在标注区按住左键拖拽一个框。"
+        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
     xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax])
-    # 归一化到 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 = run(MODEL, embedding, box_1024, H, W)  
+
 
-    # 黑白 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,  # ← 调你的推理函数
+    fn=segment,  
     inputs=BBoxAnnotator(
-        value=example,                 # 默认示例；组件里自带“上传”按钮，可以换图
+        value=example,                 
         categories=["cell", "nucleus"],
         label="upload"
     ),
@@ -134,9 +105,9 @@ if __name__ == "__main__":
     demo.queue().launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,         # 不需要 public link，HF 会自动映射
+        share=False,         
         show_error=True,
-        ssr_mode=False       # 关闭 SSR（这个是触发崩溃的常见元凶）
+        ssr_mode=False       
     )