Download ZED SDK models

download_zed_models.py

import os
import argparse
import requests
import sys

def print_help():
    help_text = """
Download ZED SDK AI models from Stereolabs.

Usage:
  python download_zed_models.py [--sdk-version {4,5,all}] [--output-dir DIR] [--print-url-only]

Options:
  --sdk-version      Specify the ZED SDK major version to filter models: '4', '5', or 'all' (default: all)
  --output-dir       Optional output directory for downloaded models
  --print-url-only   Only print the download URLs instead of downloading the files
  -h, --help         Show this help message and exit

Examples:
  python download_zed_models.py
  python download_zed_models.py --sdk-version 4
  python download_zed_models.py --sdk-version 5 --output-dir ./my_models/
  python download_zed_models.py --print-url-only
"""
    print(help_text)

def show_progress_bar(downloaded, total, bar_length=40):
    percent = float(downloaded) / total if total else 0
    arrow = '=' * int(round(percent * bar_length) - 1) + '>' if percent > 0 else ''
    spaces = ' ' * (bar_length - len(arrow))
    mb_downloaded = downloaded / 1024 / 1024
    mb_total = total / 1024 / 1024 if total else 0
    sys.stdout.write(f"\rProgress: [{arrow}{spaces}] {percent*100:.1f}% ({mb_downloaded:.2f}MB/{mb_total:.2f}MB)")
    sys.stdout.flush()

def download_files(base_url, model_list, output_folder):
    os.makedirs(output_folder, exist_ok=True)
    for model_filename in model_list:
        url_path = model_filename
        filename = model_filename
        try:
            full_url = base_url.rstrip('/') + '/' + url_path.lstrip('/')
            with requests.get(full_url, stream=True, allow_redirects=True) as response:
                final_response = response
                if response.history:
                    final_response = response.history[-1]
                if final_response.status_code == 200 or final_response.status_code == 302:
                    total_length = int(response.headers.get('content-length', 0))
                    if total_length < 1024 * 1024:
                        print(f"Error: The file {filename} is too small (less than 1MB, {total_length/1024:.2f} kB). Skipping download. Please check the model URL or try again later.")
                        continue

                    file_path = os.path.join(output_folder, filename)
                    print(f"- Downloading {filename} from {full_url} to {file_path}")
                    downloaded = 0
                    chunk_size = 8192
                    with open(file_path, 'wb') as f:
                        for chunk in response.iter_content(chunk_size=chunk_size):
                            if chunk:
                                f.write(chunk)
                                downloaded += len(chunk)
                                show_progress_bar(downloaded, total_length)
                    sys.stdout.write('\n')
                else:
                    print(f"Failed to download {filename} from {full_url} (Final Status code: {final_response.status_code})")
        except requests.exceptions.RequestException as e:
            print(f"Error occurred while downloading {filename} from {full_url}: {e}")

def print_urls(base_url, model_list):
    for model_filename in model_list:
        url_path = model_filename
        filename = model_filename
        full_url = base_url.rstrip('/') + '/' + url_path.lstrip('/')
        print(f"{filename}: {full_url}")

def main():
    print_help()

    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument(
        '--sdk-version', 
        choices=['4', '5', 'all'], 
        default='all'
    )
    parser.add_argument(
        '--output-dir',
        default=None
    )
    parser.add_argument(
        '--print-url-only',
        action='store_true',
        help='Only print the download URLs instead of downloading the files'
    )
    parser.add_argument(
        '-h', '--help',
        action='store_true'
    )
    args = parser.parse_args()

    if args.help:
        print_help()
        sys.exit(0)

    base_url = 'http://download.stereolabs.com/ai/model/'

    url_to_filename_ZEDSDK_4_X = [
        'objects_performance_3.2.model',
        'objects_medium_3.2.model',
        'objects_accurate_3.2.model',
        'skeleton_body18_3.2.model',
        'skeleton_body38_3.5.model',
        'person_head_performance_2.4.model',
        'person_head_accurate_2.4.model',
        'person_reid_1.4.model',
        'neural_depth_2.0.model',
        'neural_depth_3.6.model'
    ]

    url_to_filename_ZEDSDK_5_0 = [
        'objects_performance_3.2.model',
        'objects_medium_3.2.model',
        'objects_accurate_3.2.model',
        'skeleton_body18_3.2.model',
        'skeleton_body38_3.5.model',
        'person_head_performance_2.4.model',
        'person_head_accurate_2.4.model',
        'person_reid_1.4.model',
        'neural_depth_5.2.model',
        'neural_depth_light_5.2.model'
    ]

    if args.sdk_version == '4':
        model_list = url_to_filename_ZEDSDK_4_X
    elif args.sdk_version == '5':
        model_list = url_to_filename_ZEDSDK_5_0
    else:  # all
        # Union, then unique (preserves order from 4 then 5)
        model_list = url_to_filename_ZEDSDK_4_X + [m for m in url_to_filename_ZEDSDK_5_0 if m not in url_to_filename_ZEDSDK_4_X]

    if args.print_url_only:
        print_urls(base_url, model_list)
        sys.exit(0)

    # Output directory logic
    if args.output_dir:
        output_directory = args.output_dir
    else:
        if os.name == 'posix':
            output_directory = '/usr/local/zed/resources/'
        elif os.name == 'nt':
            output_directory = 'C:/ProgramData/Stereolabs/resources/'
        else:
            output_directory = './resources/'  # fallback for other OS

    # Ensure the output directory exists
    os.makedirs(output_directory, exist_ok=True)

    download_files(base_url, model_list, output_directory)

if __name__ == "__main__":
    main()