darwin_demo.py

import argparse
from pathlib import Path
from typing import Optional

from darwin.cli_functions import authenticate
from darwin.client import Client
from darwin.dataset.identifier import DatasetIdentifier
from darwin.dataset.split_manager import split_dataset


def run_demo(
    *,
    team_slug: Optional[str],
    dataset_slug: Optional[str] = None,
    datasets_dir: Optional[str] = None,
    api_key: Optional[str] = None,
    config_path: Optional[Path] = None,
):
    """
    Download a Darwin dataset on the file system.

    Parameters
    ----------
    team_slug : str
        Slug of the team to select
    dataset_slug : str
        This is the dataset name with everything lower-case, removed specials characters and
        spaces are replaced by dashes, e.g., `bird-species`. This string is unique within a team
    datasets_dir : Path
        Path where the client should be initialized from (aka the root path)
    api_key: str
        API key to authenticate the client
    config_path: Path
        Path to a configuration path which contains the authentication information to use

    Returns
    -------
    splits : dict
        Keys are the different splits (random, tags, ...) and values are the relative file names
    """
    # Authenticate the new KEY if available
    if api_key is not None:
        authenticate(api_key=api_key, default_team=True, datasets_dir=datasets_dir)
    # Get the client used to perform remote operations
    if config_path is not None:
        client = Client.from_config(config_path=config_path)
    else:
        client = Client.local(team_slug=team_slug)
    # Create a dataset identifier
    dataset_identifier = DatasetIdentifier(
        dataset_slug=dataset_slug, team_slug=team_slug
    )
    # Get an object representing the remote dataset
    ds = client.get_remote_dataset(dataset_identifier=dataset_identifier)
    # Download the dataset on the local file system
    ds.pull()

    # Infer the base dataset path if not specified
    if datasets_dir is None:
        datasets_dir = ds.local_path
    # Split the dataset in train/val/test
    splits = split_dataset(datasets_dir)
    # Here you can start your Machine Learning model :)
    print(splits)


if __name__ == "__main__":
    # Parse arguments
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="This script can be used to download a dataset from Darwin",
    )
    parser.add_argument(
        "--datasets-dir",
        help="Path to where the dataset will be downloaded",
        default=None,
        type=Path,
    )
    parser.add_argument(
        "--dataset-slug",
        help="Dataset slug (see Darwin documentation)",
        default=None,
        type=str,
    )
    parser.add_argument(
        "--team-slug",
        help="Team slug (see Darwin documentation)",
        default=None,
        type=str,
    )
    parser.add_argument(
        "--api-key", help="API key to authenticate the client", default=None, type=str
    )
    parser.add_argument(
        "--config-path",
        help="Path to the configuration file to authenticate the client",
        default=None,
        type=Path,
    )
    args = parser.parse_args()

    # Run the actual code
    run_demo(**args.__dict__)