1

scripts/amg_box.py

@@ -15,6 +15,7 @@ class SAMBoxSegmenter:
 
         # 初始化SAM模型
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(self.device)
         self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
         self.sam.to(device=self.device)
         self.predictor = SamPredictor(self.sam)
@@ -229,8 +230,8 @@ class SAMBoxSegmenter:
 
 if __name__ == "__main__":
     # 配置参数
-    IMAGE_PATH = "/workspace/PycharmProjects/segment-anything/scripts/ae976103-d7ec-4eed-b5d1-3e5f04d8be26.jpg"  # 替换为你的图像路径
-    SAM_CHECKPOINT = "/workspace/PycharmProjects/segment-anything/checkpoint/sam_vit_h_4b8939.pth"  # 替换为你的SAM模型路径
+    IMAGE_PATH = "ae976103-d7ec-4eed-b5d1-3e5f04d8be26.jpg"  # 替换为你的图像路径
+    SAM_CHECKPOINT = "/mnt/data/workspace/Code/aida_seg_anything/checkpoint/sam_vit_h_4b8939.pth"  # 替换为你的SAM模型路径
     MODEL_TYPE = "vit_h"  # 模型类型，与checkpoint对应
 
     # 创建并运行分割器