50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
#!/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()
|