Update app.py

app.py

@@ -13,6 +13,15 @@ import base64
 import cv2
 from typing import List, Tuple, Optional
 
+# ------------------ 安装SAM2依赖 ------------------
+# ------------------ 修复SAM2模块路径问题 ------------------
+def add_sam2_to_path():
+    """将SAM2安装目录添加到Python路径"""
+    sam2_dir = os.path.abspath("third_party/sam2")
+    if sam2_dir not in sys.path:
+        sys.path.insert(0, sam2_dir)
+    return sam2_dir
+
 # ------------------ 安装SAM2依赖 ------------------
 def install_sam2():
     """检查并安装SAM2及其依赖"""
@@ -21,37 +30,66 @@ def install_sam2():
         print("Installing SAM2...")
         os.makedirs("third_party", exist_ok=True)
         
-        # 克隆仓库
-        subprocess.run(["git", "clone", "https://github.com/facebookresearch/sam2.git", sam2_dir], check=True)
+        # 克隆仓库（添加--recursive以防有子模块）
+        subprocess.run([
+            "git", "clone", 
+            "--recursive", 
+            "https://github.com/facebookresearch/sam2.git", 
+            sam2_dir
+        ], check=True)
         
         # 切换到sam2目录安装
         original_dir = os.getcwd()
         try:
             os.chdir(sam2_dir)
             
-            # 安装依赖
-            subprocess.run(["python","-m","pip", "install", "-e", "."], check=True)
+            # 先安装核心依赖
+            subprocess.run(["pip", "install", "-r", "requirements.txt"], check=True)
+            
+            # 以可编辑模式安装SAM2
+            subprocess.run(["pip", "install", "-e", "."], check=True)
             
             # 切换到checkpoints目录下载模型
             os.chdir("checkpoints")
             if not os.path.exists("download_ckpts.sh"):
-                # 如果脚本不存在，可能是权限问题，尝试直接下载
-                subprocess.run(["wget", "https://raw.githubusercontent.com/facebookresearch/sam2/main/checkpoints/download_ckpts.sh"], check=True)
+                subprocess.run([
+                    "wget", 
+                    "https://raw.githubusercontent.com/facebookresearch/sam2/main/checkpoints/download_ckpts.sh"
+                ], check=True)
                 subprocess.run(["chmod", "+x", "download_ckpts.sh"], check=True)
             
-            subprocess.run(["./download_ckpts.sh"], check=True)
+            # 下载检查点（添加超时和重试）
+            result = subprocess.run(["./download_ckpts.sh"], check=False)
+            if result.returncode != 0:
+                print("Warning: Checkpoint download failed, trying alternative method...")
+                subprocess.run([
+                    "wget",
+                    "https://dl.fbaipublicfiles.com/sam2/sam2.1_hiera_large.pt",
+                    "-O", "sam2.1_hiera_large.pt"
+                ], check=True)
             
+        except Exception as e:
+            print(f"Error during SAM2 installation: {str(e)}")
+            raise
         finally:
-            # 确保切换回原始目录
             os.chdir(original_dir)
         
-        print("SAM2 installed successfully!")
+        print("✅ SAM2 installed successfully!")
     else:
-        print("SAM2 already installed.")
+        print("SAM2 already exists, skipping installation.")
 
-# 确保安装SAM2
+
+    # 1. 安装SAM2
 install_sam2()
 
+# 2. 确保路径正确
+sam2_dir = add_sam2_to_path()
+
+# 3. 现在可以安全导入SAM2模块
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+print("🎉 SAM2 modules imported successfully!")
+
 # ------------------ 初始化模型 ------------------
 # 使用相对路径
 MODEL_PATH = "geshang/Seg-R1-COD"