File size: 3,329 Bytes
0e83290
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import requests

# URL файла для скачивания
url = 'https://github.com/hkchengrex/XMem/releases/download/v1.0/XMem.pth'
# Имя файла, под которым он будет сохранен
file_name = 'XMem.pth'
path = 'checkpoints'
# Базовые URL-адреса
BASE_URL_PT_1 = "https://dl.fbaipublicfiles.com/segment_anything_2/092824"
BASE_URL_YAML_1 = (
    "https://raw.githubusercontent.com/facebookresearch/sam2/main/sam2/configs/sam2.1"
)


# Словарь для чекпоинтов (.pt)
CHECKPOINTS_PT_2_1 = {
    f"{path}/sam2.1_hiera_tiny.pt": f"{BASE_URL_PT_1}/sam2.1_hiera_tiny.pt",
    f"{path}/sam2.1_hiera_small.pt": f"{BASE_URL_PT_1}/sam2.1_hiera_small.pt",
    f"{path}/sam2.1_hiera_base_plus.pt": f"{BASE_URL_PT_1}/sam2.1_hiera_base_plus.pt",
    f"{path}/sam2.1_hiera_large.pt": f"{BASE_URL_PT_1}/sam2.1_hiera_large.pt",
}

# Словарь для конфигурационных файлов (.yaml)
# CHECKPOINTS_YAML_1 = {
#     f"{path}/sam2.1_hiera_t.yaml": f"{BASE_URL_YAML_1}/sam2.1_hiera_t.yaml",
#     f"{path}/sam2.1_hiera_s.yaml": f"{BASE_URL_YAML_1}/sam2.1_hiera_s.yaml",
#     f"{path}/sam2.1_hiera_b+.yaml": f"{BASE_URL_YAML_1}/sam2.1_hiera_b+.yaml",
#     f"{path}/sam2.1_hiera_l.yaml": f"{BASE_URL_YAML_1}/sam2.1_hiera_l.yaml",
# }

BASE_URL_PT = "https://dl.fbaipublicfiles.com/segment_anything_2/072824"
BASE_URL_YAML = "https://raw.githubusercontent.com/Segment-Anything/segment-anything-2/main/sam2_configs"
# https://raw.githubusercontent.com/facebookresearch/sam2/main/sam2/configs/sam2

CHECKPOINTS_PT_2 = {
    f"{path}/sam2_hiera_tiny.pt": f"{BASE_URL_PT}/sam2_hiera_tiny.pt",
    f"{path}/sam2_hiera_small.pt": f"{BASE_URL_PT}/sam2_hiera_small.pt",
    f"{path}/sam2_hiera_base_plus.pt": f"{BASE_URL_PT}/sam2_hiera_base_plus.pt",
    f"{path}/sam2_hiera_large.pt": f"{BASE_URL_PT}/sam2_hiera_large.pt",
}

# Словарь для конфигурационных файлов (.yaml)
# CHECKPOINTS_YAML = {
#     f"{path}/sam2_hiera_t.yaml": f"{BASE_URL_YAML}/sam2_hiera_t.yaml",
#     f"{path}/sam2_hiera_s.yaml": f"{BASE_URL_YAML}/sam2_hiera_s.yaml",
#     f"{path}/sam2_hiera_b+.yaml": f"{BASE_URL_YAML}/sam2_hiera_b+.yaml",
#     f"{path}/sam2_hiera_l.yaml": f"{BASE_URL_YAML}/sam2_hiera_l.yaml",
# }


def download_checkpoint(url, filename):
    print(f"Downloading {filename}...")
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()  # Проверка на ошибки HTTP

        with open(filename, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)

        print(f"{filename} downloaded successfully.")
    except requests.exceptions.RequestException as e:
        print(f"Failed to download {filename} from {url}. Error: {e}")
        exit(1)


if __name__ == "__main__":
    download_checkpoint(url, f'{path}/{file_name}')
    # Сначала скачиваем .pt файлы
    for filename, url in CHECKPOINTS_PT_2.items():
        download_checkpoint(url, filename)

    for filename, url in CHECKPOINTS_PT_2_1.items():
        download_checkpoint(url, filename)

    print("All checkpoints and configuration files are downloaded successfully.")