Update inference_seg.py

inference_seg.py

@@ -1,24 +1,66 @@
+# import torch
+# import numpy as np
+# from huggingface_hub import hf_hub_download
+# from segmentation import SegmentationModule  
+
+# MODEL = None
+# DEVICE = torch.device("cpu")
+
+# def load_model(use_box=False):
+#     global MODEL, DEVICE
+#     MODEL = CountingModule(use_box=use_box)
+
+#     ckpt_path = hf_hub_download(
+#         repo_id="Shengxiao0709/cellsegmodel",
+#         filename="microscopy_matching_seg.pth",
+#         token=None,
+#         force_download=False
+#     )
+#     MODEL.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
+#     MODEL.eval()
+#     DEVICE = torch.device("cpu")
+#     return MODEL, DEVICE
+
+
+# @torch.no_grad()
+# def run(model, img_path, box=None, device="cpu"):
+#     output = model(img_path, box=box)
+#     mask = output["pred"]
+#     mask = (mask > 0).astype(np.uint8)
+#     return mask
+import os
 import torch
 import numpy as np
 from huggingface_hub import hf_hub_download
-from segmentation import SegmentationModule  
+from segmentation import SegmentationModule
 
 MODEL = None
 DEVICE = torch.device("cpu")
 
 def load_model(use_box=False):
     global MODEL, DEVICE
-    MODEL = CountingModule(use_box=use_box)
+
+    # === 优化1: 使用 /data 缓存模型，避免写入 .cache ===
+    cache_dir = "/data/cellseg_model_cache"
+    os.makedirs(cache_dir, exist_ok=True)
 
     ckpt_path = hf_hub_download(
         repo_id="Shengxiao0709/cellsegmodel",
         filename="microscopy_matching_seg.pth",
         token=None,
-        force_download=False
+        local_dir=cache_dir,              # ✅ 下载到 /data
+        local_dir_use_symlinks=False,     # ✅ 避免软链接问题
+        force_download=False              # ✅ 已存在时不重复下载
     )
-    MODEL.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
+
+    # === 优化2: 加载模型 ===
+    MODEL = SegmentationModule(use_box=use_box)
+    state_dict = torch.load(ckpt_path, map_location="cpu")
+    MODEL.load_state_dict(state_dict, strict=False)
     MODEL.eval()
+
     DEVICE = torch.device("cpu")
+    print(f"✅ Model loaded from {ckpt_path}")
     return MODEL, DEVICE