gbreadman13/sam2-api

1#!/usr/bin/env python32"""3Скрипт для скачивания модели SAM2.4Блин, Facebook не может нормально в pip packaging, поэтому качаем руками.5"""6 7import os8import urllib.request9import sys10 11# Директория для чекпоинтов12CHECKPOINT_DIR = "checkpoints"13os.makedirs(CHECKPOINT_DIR, exist_ok=True)14 15# Модели на выбор16MODELS = {17    "tiny": {18        "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt",19        "filename": "sam2.1_hiera_tiny.pt",20        "size": "~39MB"21    },22    "small": {23        "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt",24        "filename": "sam2.1_hiera_small.pt",25        "size": "~46MB"26    },27    "base_plus": {28        "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt",29        "filename": "sam2.1_hiera_base_plus.pt",30        "size": "~81MB"31    },32    "large": {33        "url": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt",34        "filename": "sam2.1_hiera_large.pt",35        "size": "~224MB"36    }37}38 39def download_model(model_name="tiny"):40    """Качает модель, показывает прогресс"""41    if model_name not in MODELS:42        print(f"Неизвестная модель: {model_name}")43        print(f"Доступные: {', '.join(MODELS.keys())}")44        sys.exit(1)45    46    model_info = MODELS[model_name]47    filepath = os.path.join(CHECKPOINT_DIR, model_info["filename"])48    49    if os.path.exists(filepath):50        print(f"Модель уже скачана: {filepath}")51        return filepath52    53    print(f"Качаю {model_name} модель ({model_info['size']})...")54    print(f"URL: {model_info['url']}")55    56    def progress_hook(block_num, block_size, total_size):57        downloaded = block_num * block_size58        if total_size > 0:59            percent = min(100, downloaded * 100 / total_size)60            sys.stdout.write(f"\rПрогресс: {percent:.1f}%")61            sys.stdout.flush()62    63    try:64        urllib.request.urlretrieve(65            model_info["url"],66            filepath,67            reporthook=progress_hook68        )69        print(f"\n✓ Модель скачана: {filepath}")70        return filepath71    except Exception as e:72        print(f"\n✗ Ошибка при скачивании: {e}")73        if os.path.exists(filepath):74            os.remove(filepath)75        sys.exit(1)76 77if __name__ == "__main__":78    model_name = sys.argv[1] if len(sys.argv) > 1 else "tiny"79    download_model(model_name)80 81