Download Pretrained Models¶

flyvis_cli.download_pretrained_models ¶

calculate_sha256 ¶

calculate_sha256(file_path)

Calculate the SHA256 checksum of a file.

Parameters:

Name	Type	Description	Default
`file_path`	`str`	Path to the file.	required

Returns:

Type	Description
`str`	The SHA256 checksum as a hexadecimal string.

Source code in flyvis_cli/download_pretrained_models.py

def calculate_sha256(file_path: str) -> str:
    """
    Calculate the SHA256 checksum of a file.

    Args:
        file_path: Path to the file.

    Returns:
        The SHA256 checksum as a hexadecimal string.
    """
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

download_and_unpack_files ¶

download_and_unpack_files(folder_id, destination_dir, api_key, skip_large_files=False)

Download and unpack files from a Google Drive folder.

Parameters:

Name	Type	Description	Default
`folder_id`	`str`	The ID of the Google Drive folder.	required
`destination_dir`	`str`	The local directory to save and unpack the files.	required
`api_key`	`str`	The Google Drive API key.	required
`skip_large_files`	`bool`	Whether to skip downloading large files.	`False`

Note

This function creates the destination directory if it doesn’t exist, downloads ZIP files from the specified Google Drive folder, checks their checksums, and unpacks them.

Source code in flyvis_cli/download_pretrained_models.py

def download_and_unpack_files(
    folder_id: str, destination_dir: str, api_key: str, skip_large_files: bool = False
) -> None:
    """
    Download and unpack files from a Google Drive folder.

    Args:
        folder_id: The ID of the Google Drive folder.
        destination_dir: The local directory to save and unpack the files.
        api_key: The Google Drive API key.
        skip_large_files: Whether to skip downloading large files.

    Note:
        This function creates the destination directory if it doesn't exist,
        downloads ZIP files from the specified Google Drive folder, checks their
        checksums, and unpacks them.
    """
    if not os.path.exists(destination_dir):
        os.makedirs(destination_dir)

    service = build("drive", "v3", developerKey=api_key)

    query = f"'{folder_id}' in parents and mimeType='application/zip'"
    results = service.files().list(q=query, fields="files(id, name)").execute()
    items = results.get("files", [])

    if not items:
        print("No files found.")
        return

    for item in items:
        file_id = item["id"]
        file_name = item["name"]

        if skip_large_files and large_files.get(file_name, False):
            print(f"Skipping {file_name}.")
            continue

        request = service.files().get_media(fileId=file_id)
        file_path = os.path.join(destination_dir, file_name)
        print(f"Downloading {file_name} to {file_path}.")

        with io.FileIO(file_path, "wb") as fh:
            downloader = MediaIoBaseDownload(fh, request)
            done = False
            while not done:
                status, done = downloader.next_chunk()
                print(f"Progress {file_name}: {int(status.progress() * 100)}%", end="\r")
            print()

        checksum = calculate_sha256(file_path)
        if checksum != checksums[file_name]:
            print(
                f"Checksum mismatch for {file_name}. Expected {checksums[file_name]}, "
                f"got {checksum}."
            )
            return
        print(f"Checksum OK for {file_name}.")

        with zipfile.ZipFile(file_path, "r") as zip_ref:
            zip_ref.extractall(destination_dir)
        print(f"Unpacked {file_name}.")

main ¶

main()

Main function to handle command line arguments and initiate the download process.

Source code in flyvis_cli/download_pretrained_models.py

def main() -> None:
    """
    Main function to handle command line arguments and initiate the download process.
    """
    arg_parser = argparse.ArgumentParser(
        description="Download pretrained models and UMAP clustering results. "
        "This script downloads two ZIP files from Google Drive:\n"
        "1. results_pretrained_models.zip: Contains pretrained neural networks.\n"
        "2. results_umap_and_clustering.zip: Contains UMAP embeddings and "
        "clustering.\n"
        "The files are downloaded and unpacked to the 'data' directory in the "
        "project root.",
        formatter_class=argparse.RawTextHelpFormatter,
        usage=(
            "\nflyvis download-pretrained [-h] [--skip_large_files]\n"
            "       or\n"
            "%(prog)s [-h] [--skip_large_files]\n"
        ),
    )
    arg_parser.add_argument(
        "--skip_large_files",
        action="store_true",
        help="Skip downloading large files. If set, only 'results_pretrained_models.zip' "
        "will be downloaded, and 'results_umap_and_clustering.zip' will be skipped.",
    )
    args, _ = arg_parser.parse_known_intermixed_args()

    folder_id = "15_8gPaVVJV6wGAspwkrdN8r-2NxMY4Kj"
    api_key = "AIzaSyDOy2_N7B5UjxKy5Xxeyd9WZfnDWzQ4-54"

    download_and_unpack_files(folder_id, str(root_dir), api_key, args.skip_large_files)

usage:
flyvis download-pretrained [-h] [--skip_large_files]
       or
download_pretrained_models.py [-h] [--skip_large_files]

Download pretrained models and UMAP clustering results. This script downloads two ZIP files from Google Drive:
1. results_pretrained_models.zip: Contains pretrained neural networks.
2. results_umap_and_clustering.zip: Contains UMAP embeddings and clustering.
The files are downloaded and unpacked to the 'data' directory in the project root.

options:
  -h, --help          show this help message and exit
  --skip_large_files  Skip downloading large files. If set, only 'results_pretrained_models.zip' will be downloaded, and 'results_umap_and_clustering.zip' will be skipped.