#!/usr/bin/env python3
"""
SAM 2 模型权重下载脚本
运行: python download_sam2.py
"""
import os
import urllib.request
import sys

MODEL_DIR = "/home/wkmgc/Desktop/Seg_Server/models"
os.makedirs(MODEL_DIR, exist_ok=True)

# SAM 2 模型权重 (Meta AI 官方)
MODELS = {
    "sam2_hiera_tiny.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt",
    "sam2_hiera_small.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
    "sam2_hiera_base_plus.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt",
    "sam2_hiera_large.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt",
}

def download_file(url: str, dest: str):
    if os.path.exists(dest):
        print(f"[跳过] {os.path.basename(dest)} 已存在")
        return
    print(f"[下载] {url}")
    print(f"  → {dest}")
    try:
        urllib.request.urlretrieve(url, dest)
        size_mb = os.path.getsize(dest) / 1024 / 1024
        print(f"  ✓ 完成 ({size_mb:.1f} MB)")
    except Exception as e:
        print(f"  ✗ 失败: {e}", file=sys.stderr)
        if os.path.exists(dest):
            os.remove(dest)

def main():
    print("=" * 50)
    print("SAM 2 模型权重下载")
    print("=" * 50)
    for name, url in MODELS.items():
        dest = os.path.join(MODEL_DIR, name)
        download_file(url, dest)
    print("=" * 50)
    print("全部完成！")
    print(f"模型目录: {MODEL_DIR}")
    print("=" * 50)

if __name__ == "__main__":
    main()