__init__.py

"""Nucleus Python SDK. """

__all__ = [
    "AsyncJob",
    "EmbeddingsExportJob",
    "BoxAnnotation",
    "DeduplicationResult",
    "DeduplicationStats",
    "BoxPrediction",
    "CameraParams",
    "CategoryAnnotation",
    "CategoryPrediction",
    "CuboidAnnotation",
    "CuboidPrediction",
    "Dataset",
    "DatasetInfo",
    "DatasetItem",
    "DatasetItemRetrievalError",
    "Frame",
    "Keypoint",
    "KeypointsAnnotation",
    "KeypointsPrediction",
    "LidarPoint",
    "LidarScene",
    "LineAnnotation",
    "LinePrediction",
    "Model",
    "ModelCreationError",
    # "MultiCategoryAnnotation", # coming soon!
    "NotFoundError",
    "NucleusAPIError",
    "NucleusClient",
    "Point",
    "Point3D",
    "PolygonAnnotation",
    "PolygonPrediction",
    "Quaternion",
    "SceneCategoryAnnotation",
    "SceneCategoryPrediction",
    "Segment",
    "SegmentationAnnotation",
    "SegmentationPrediction",
    "Slice",
    "VideoScene",
]

import datetime
import os
import warnings
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import requests
import tqdm

if TYPE_CHECKING:
    # Backwards compatibility is even uglier with mypy
    from pydantic.v1 import parse_obj_as
else:
    try:
        # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
        from pydantic.v1 import parse_obj_as
    except ImportError:
        from pydantic import parse_obj_as

from nucleus.url_utils import sanitize_string_args

from . import metrics
from .annotation import (
    BoxAnnotation,
    CategoryAnnotation,
    CuboidAnnotation,
    Keypoint,
    KeypointsAnnotation,
    LidarPoint,
    LineAnnotation,
    MultiCategoryAnnotation,
    Point,
    Point3D,
    PolygonAnnotation,
    SceneCategoryAnnotation,
    Segment,
    SegmentationAnnotation,
)
from .async_job import AsyncJob, EmbeddingsExportJob
from .async_utils import make_multiple_requests_concurrently
from .camera_params import CameraParams
from .connection import Connection
from .constants import (
    ANNOTATION_METADATA_SCHEMA_KEY,
    ANNOTATIONS_IGNORED_KEY,
    ANNOTATIONS_PROCESSED_KEY,
    AUTOTAGS_KEY,
    DATASET_ID_KEY,
    DATASET_IS_SCENE_KEY,
    DATASET_PRIVACY_MODE_KEY,
    DEFAULT_NETWORK_TIMEOUT_SEC,
    EMBEDDING_DIMENSION_KEY,
    EMBEDDINGS_URL_KEY,
    ERROR_ITEMS,
    ERROR_PAYLOAD,
    ERRORS_KEY,
    GLOB_SIZE_THRESHOLD_CHECK,
    I_KEY,
    IMAGE_KEY,
    IMAGE_URL_KEY,
    INDEX_CONTINUOUS_ENABLE_KEY,
    ITEM_METADATA_SCHEMA_KEY,
    ITEMS_KEY,
    JOB_CREATION_TIME_KEY,
    JOB_ID_KEY,
    JOB_LAST_KNOWN_STATUS_KEY,
    JOB_TYPE_KEY,
    KEEP_HISTORY_KEY,
    MESSAGE_KEY,
    MODEL_RUN_ID_KEY,
    MODEL_TAGS_KEY,
    MODEL_TRAINED_SLICE_IDS_KEY,
    NAME_KEY,
    NUCLEUS_ENDPOINT,
    POINTS_KEY,
    PREDICTIONS_IGNORED_KEY,
    PREDICTIONS_PROCESSED_KEY,
    REFERENCE_IDS_KEY,
    SLICE_ID_KEY,
    SLICE_TAGS_KEY,
    STATUS_CODE_KEY,
    UPDATE_KEY,
)
from .data_transfer_object.dataset_details import DatasetDetails
from .data_transfer_object.dataset_info import DatasetInfo
from .data_transfer_object.job_status import JobInfoRequestPayload
from .dataset import Dataset
from .dataset_item import DatasetItem
from .deduplication import DeduplicationResult, DeduplicationStats
from .deprecation_warning import deprecated
from .errors import (
    DatasetItemRetrievalError,
    ModelCreationError,
    ModelRunCreationError,
    NoAPIKey,
    NotFoundError,
    NucleusAPIError,
)
from .job import CustomerJobTypes
from .model import Model
from .model_run import ModelRun
from .payload_constructor import (
    construct_annotation_payload,
    construct_append_payload,
    construct_box_predictions_payload,
    construct_model_creation_payload,
    construct_segmentation_payload,
)
from .prediction import (
    BoxPrediction,
    CategoryPrediction,
    CuboidPrediction,
    KeypointsPrediction,
    LinePrediction,
    PolygonPrediction,
    SceneCategoryPrediction,
    SegmentationPrediction,
)
from .quaternion import Quaternion
from .retry_strategy import RetryStrategy
from .scene import Frame, LidarScene, VideoScene
from .slice import Slice
from .upload_response import UploadResponse
from .utils import create_items_from_folder_crawl
from .validate import Validate

# pylint: disable=E1101
# TODO: refactor to reduce this file to under 1000 lines.
# pylint: disable=C0302


class NucleusClient:
    """Client to interact with the Nucleus API via Python SDK.

    Parameters:
        api_key: One of ``api_key`` or ``limited_access_key`` must be provided; you cannot pass
          both. For standard Scale API key authentication, pass the key here. Follow `this guide
          <https://scale.com/docs/api-reference/authentication>`_ to retrieve API keys. If you omit
          this argument and are not using ``limited_access_key``, the SDK falls back to the
          ``NUCLEUS_API_KEY`` environment variable.
        limited_access_key: Nucleus-only API key for scoped access. Reach out to your Scale
          representative to obtain a limited access key.
        use_notebook: Whether the client is being used in a notebook (toggles tqdm
          style). Default is ``False``.
        endpoint: Base URL of the API. Default is Nucleus's current production API.

    .. note::

        You must provide **either** a standard Scale API key (``api_key``, or
        ``NUCLEUS_API_KEY`` in the environment) **or** a Nucleus-only key
        (``limited_access_key``), never both. Passing both arguments, or setting
        the environment variable ``NUCLEUS_API_KEY`` while also passing
        ``limited_access_key``, will raise an error.

    Example::

        # Using a basic auth key
        import nucleus
        client = nucleus.NucleusClient(api_key="YOUR_API_KEY", ...)

        # Using only a limited access key (no Basic Auth)
        import nucleus
        client = nucleus.NucleusClient(limited_access_key="YOUR_LIMITED_KEY", ...)
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        use_notebook: bool = False,
        endpoint: Optional[str] = None,
        limited_access_key: Optional[str] = None,
    ):
        effective_basic_key = (
            api_key if api_key else os.environ.get("NUCLEUS_API_KEY")
        )
        if limited_access_key and effective_basic_key:
            raise ValueError(
                "Cannot provide both 'api_key' and 'limited_access_key'. "
                "Use 'api_key' for standard Scale API key authentication, "
                "or 'limited_access_key' for Nucleus-only access, but not both."
            )
        # Allow usage with only a limited access key
        if api_key is None and limited_access_key:
            self.api_key = None
        else:
            self.api_key = self._set_api_key(api_key)
        self.tqdm_bar = tqdm.tqdm
        if endpoint is None:
            self.endpoint = os.environ.get(
                "NUCLEUS_ENDPOINT", NUCLEUS_ENDPOINT
            )
        else:
            self.endpoint = endpoint
        self._use_notebook = use_notebook
        if use_notebook:
            import tqdm.notebook as tqdm_notebook

            self.tqdm_bar = tqdm_notebook.tqdm
        self.extra_headers: Dict[str, str] = {}
        if limited_access_key:
            self.extra_headers["x-limited-access-key"] = limited_access_key
        self.connection = Connection(self.api_key, self.endpoint, extra_headers=self.extra_headers)
        self.validate = Validate(self.api_key, self.endpoint, extra_headers=self.extra_headers)

    def __repr__(self):
        return f"NucleusClient(api_key='{self.api_key}', use_notebook={self._use_notebook}, endpoint='{self.endpoint}')"

    def __eq__(self, other):
        if self.api_key == other.api_key:
            if self._use_notebook == other._use_notebook:
                return True
        return False

    @property
    def datasets(self) -> List[Dataset]:
        """List all Datasets

        Returns:
            List of all datasets accessible to user
        """
        response = self.make_request({}, "dataset/details", requests.get)
        dataset_details = (
            parse_obj_as(  # pylint: disable=used-before-assignment
                List[DatasetDetails], response
            )
        )
        return [
            Dataset(d.id, client=self, name=d.name) for d in dataset_details
        ]

    @property
    def models(self) -> List[Model]:
        # TODO: implement for Dataset, scoped just to associated models
        """Fetches all of your Nucleus models.

        Returns:
            List[:class:`Model`]: List of models associated with the client API key.
        """
        model_objects = self.make_request({}, "models/", requests.get)

        return [
            Model(
                model_id=model["id"],
                name=model["name"],
                reference_id=model["ref_id"],
                metadata=model["metadata"] or None,
                client=self,
                tags=model.get(MODEL_TAGS_KEY, []),
                trained_slice_ids=model.get(MODEL_TRAINED_SLICE_IDS_KEY, None),
            )
            for model in model_objects["models"]
        ]

    @property
    def jobs(
        self,
    ) -> List[AsyncJob]:
        """Lists all jobs, see NucleusClient.list_jobs(...) for advanced options

        Returns:
            List of all AsyncJobs
        """
        return self.list_jobs()

    @property
    def slices(self) -> List[Slice]:
        response = self.make_request({}, "slice/", requests.get)
        slices = [Slice.from_request(info, self) for info in response]
        return slices

    @deprecated(msg="Use the NucleusClient.models property in the future.")
    def list_models(self) -> List[Model]:
        return self.models

    @deprecated(msg="Use the NucleusClient.datasets property in the future.")
    def list_datasets(self) -> Dict[str, Union[str, List[str]]]:
        return self.make_request({}, "dataset/", requests.get)

    def list_jobs(
        self,
        show_completed: bool = False,
        from_date: Optional[Union[str, datetime.datetime]] = None,
        to_date: Optional[Union[str, datetime.datetime]] = None,
        job_types: Optional[List[CustomerJobTypes]] = None,
        limit: Optional[int] = None,
        dataset_id: Optional[str] = None,
        date_limit: Optional[str] = None,
    ) -> List[AsyncJob]:
        """Fetches all of your running jobs in Nucleus.

        Parameters:
            show_completed: Whether to include jobs with Completed status.
            from_date: Beginning of date range filter.
            to_date: End of date range filter.
            job_types: Filter on set of job types. If None, fetch all types.
            limit: Number of results to fetch, max 50,000.
            dataset_id: Filter on a particular dataset.
            date_limit: Deprecated, do not use.

        Returns:
            List[:class:`AsyncJob`]: List of running asynchronous jobs
            associated with the client API key.
        """

        if date_limit is not None:
            warnings.warn(
                "Argument `date_limit` is no longer supported. Consider using the `from_date` and `to_date` args."
            )

        payload = JobInfoRequestPayload(
            dataset_id=dataset_id,
            show_completed=show_completed,
            from_date=from_date,
            to_date=to_date,
            limit=limit,
            job_types=job_types,
        ).dict()

        job_objects = self.make_request(payload, "jobs/", requests.post)
        return [
            AsyncJob(
                job_id=job[JOB_ID_KEY],
                job_last_known_status=job[JOB_LAST_KNOWN_STATUS_KEY],
                job_type=job[JOB_TYPE_KEY],
                job_creation_time=job[JOB_CREATION_TIME_KEY],
                client=self,
            )
            for job in job_objects
        ]

    @deprecated(msg="Prefer using Dataset.items")
    def get_dataset_items(self, dataset_id) -> List[DatasetItem]:
        dataset = self.get_dataset(dataset_id)
        return dataset.items

    def get_dataset(self, dataset_id: str) -> Dataset:
        """Fetches a dataset by its ID.

        Parameters:
            dataset_id: The ID of the dataset to fetch.

        Returns:
            :class:`Dataset`: The Nucleus dataset as an object.
        """
        return Dataset(dataset_id, self)

    def get_job(self, job_id: str) -> AsyncJob:
        """Fetches a job by its ID.

        Parameters:
            job_id: The ID of the job to fetch.

        Returns:
            :class:`AsyncJob`: The Nucleus async job as an object.
        """
        payload = self.make_request(
            payload={},
            route=f"job/{job_id}/info",
            requests_command=requests.get,
        )
        return AsyncJob.from_json(payload=payload, client=self)

    def get_model(
        self,
        model_id: Optional[str] = None,
        model_run_id: Optional[str] = None,
    ) -> Model:
        """Fetches a model by its ID.

        Parameters:
            model_id: You can pass either a model ID (starts with ``prj_``) or a
                model run ID (starts with ``run_``). Retrieved via :meth:`list_models`
                or from a Nucleus dashboard URL.
            model_run_id: You can pass either a model ID (starts with ``prj_``), or a
                model run ID (starts with ``run_``). This can be retrieved via
                :meth:`list_models` or a Nucleus dashboard URL. Model run IDs result
                from the application of a model to a dataset. In the future, we plan
                to hide ``model_run_ids`` fully from users.

        Returns:
            :class:`Model`: The Nucleus model as an object.
        """
        if model_id is None and model_run_id is None:
            raise ValueError("Must pass either a model_id or a model_run_id")
        if model_id is not None and model_run_id is not None:
            raise ValueError("Must pass either a model_id or a model_run_id")

        model_or_model_run_id = (
            model_id if model_id is not None else model_run_id
        )

        payload = self.make_request(
            payload={},
            route=f"model/{model_or_model_run_id}",
            requests_command=requests.get,
        )
        return Model.from_json(payload=payload, client=self)

    @deprecated(
        "Model runs have been deprecated and will be removed. Use a Model instead"
    )
    def get_model_run(self, model_run_id: str, dataset_id: str) -> ModelRun:
        return ModelRun(model_run_id, dataset_id, self)

    @deprecated(
        "Model runs have been deprecated and will be removed. Use a Model instead"
    )
    def delete_model_run(self, model_run_id: str):
        return self.make_request(
            {}, f"modelRun/{model_run_id}", requests.delete
        )

    def create_dataset_from_project(
        self,
        project_id: str,
        last_n_tasks: Optional[int] = None,
        name: Optional[str] = None,
    ) -> Dataset:
        """Create a new dataset from an existing Scale or Rapid project.

        If you already have Annotation, SegmentAnnotation, VideoAnnotation,
        Categorization, PolygonAnnotation, ImageAnnotation, DocumentTranscription,
        LidarLinking, LidarAnnotation, or VideoboxAnnotation projects with Scale,
        use this endpoint to import your project directly into Nucleus.

        This endpoint is asynchronous because there can be delays when the
        number of tasks is larger than 1000. As a result, the endpoint returns
        an instance of :class:`AsyncJob`.

        Parameters:
            project_id: The ID of the Scale/Rapid project (retrievable from URL).
            last_n_tasks: If supplied, only pull in this number of the most recent
              tasks. By default the endpoint will pull in all eligible tasks.
            name: The name for your new Nucleus dataset. By default the endpoint
              will use the project's name.

        Returns:
            :class:`Dataset`: The newly created Nucleus dataset as an object.
        """
        payload = {"project_id": project_id}
        if last_n_tasks:
            payload["last_n_tasks"] = str(last_n_tasks)
        if name:
            payload["name"] = name
        response = self.make_request(payload, "dataset/create_from_project")
        return Dataset(response[DATASET_ID_KEY], self)

    def create_dataset(
        self,
        name: str,
        is_scene: Optional[bool] = None,
        use_privacy_mode: bool = False,
        item_metadata_schema: Optional[Dict] = None,
        annotation_metadata_schema: Optional[Dict] = None,
    ) -> Dataset:
        """
        Creates a new, empty dataset.

        Make sure that the dataset is created for the data type you would like to support.
        Be sure to set the ``is_scene`` parameter correctly.

        Parameters:
            name: A human-readable name for the dataset.
            is_scene: Whether the dataset contains strictly :class:`scenes
              <LidarScene>` or :class:`items <DatasetItem>`. This value is immutable.
              Default is False (dataset of items).
            use_privacy_mode: Whether the images of this dataset should be uploaded to Scale. If set to True,
              customer will have to adjust their file access policy with Scale.
            item_metadata_schema: Dict defining item-level metadata schema, structured as::

                {
                    "field_name": {
                        "type": "category" | "number" | "text" | "json"
                        "choices": List[str] | None
                        "description": str | None
                    },
                    ...
                }

            annotation_metadata_schema: Dict defining annotation-level metadata schema.
              Same format as ``item_metadata_schema``.

        Returns:
            :class:`Dataset`: The newly created Nucleus dataset as an object.
        """
        if is_scene is None:
            warnings.warn(
                "The default create_dataset('dataset_name', ...) method without the is_scene parameter will be "
                "deprecated soon in favor of providing the is_scene parameter explicitly. "
                "Please make sure to create a dataset with either create_dataset('dataset_name', is_scene=False, ...) "
                "to upload DatasetItems or create_dataset('dataset_name', is_scene=True, ...) to upload LidarScenes.",
                DeprecationWarning,
            )
            is_scene = False
        response = self.make_request(
            {
                NAME_KEY: name,
                DATASET_IS_SCENE_KEY: is_scene,
                DATASET_PRIVACY_MODE_KEY: use_privacy_mode,
                ANNOTATION_METADATA_SCHEMA_KEY: annotation_metadata_schema,
                ITEM_METADATA_SCHEMA_KEY: item_metadata_schema,
            },
            "dataset/create",
        )
        return Dataset(
            response[DATASET_ID_KEY],
            self,
            name=name,
            is_scene=is_scene,
            use_privacy_mode=use_privacy_mode,
        )

    def delete_dataset(self, dataset_id: str) -> dict:
        """
        Deletes a dataset by ID.

        All items, annotations, and predictions associated with the dataset will
        be deleted as well. Note that if this dataset is linked to a Scale or Rapid
        labeling project, the project itself will not be deleted.

        Parameters:
            dataset_id: The ID of the dataset to delete.

        Returns:
            Payload to indicate deletion invocation.
        """
        return self.make_request({}, f"dataset/{dataset_id}", requests.delete)

    @deprecated("Use Dataset.delete_item instead.")
    def delete_dataset_item(self, dataset_id: str, reference_id) -> dict:
        dataset = self.get_dataset(dataset_id)
        return dataset.delete_item(reference_id)

    @deprecated("Use Dataset.append instead.")
    def populate_dataset(
        self,
        dataset_id: str,
        dataset_items: List[DatasetItem],
        batch_size: int = 20,
        update: bool = False,
    ):
        dataset = self.get_dataset(dataset_id)
        return dataset.append(
            dataset_items, batch_size=batch_size, update=update
        )

    @deprecated(msg="Use Dataset.ingest_tasks instead")
    def ingest_tasks(self, dataset_id: str, payload: dict):
        dataset = self.get_dataset(dataset_id)
        return dataset.ingest_tasks(payload["tasks"])

    @deprecated(msg="Use client.create_model instead.")
    def add_model(
        self, name: str, reference_id: str, metadata: Optional[Dict] = None
    ) -> Model:
        return self.create_model(name, reference_id, metadata)

    def create_model(
        self,
        name: str,
        reference_id: str,
        metadata: Optional[Dict] = None,
        bundle_name: Optional[str] = None,
        tags: Optional[List[str]] = None,
        trained_slice_ids: Optional[List[str]] = None,
    ) -> Model:
        """Adds a :class:`Model` to Nucleus.

        Parameters:
            name: A human-readable name for the model.
            reference_id: Unique, user-controlled ID for the model. This can be
              used, for example, to link to an external storage of models which
              may have its own id scheme.
            metadata: An arbitrary dictionary of additional data about this model
              that can be stored and retrieved. For example, you can store information
              about the hyperparameters used in training this model.
            bundle_name: Optional name of bundle attached to this model
            tags: Optional list of tags to attach to this model

        Returns:
            :class:`Model`: The newly created model as an object.
        """
        response = self.make_request(
            construct_model_creation_payload(
                name,
                reference_id,
                metadata,
                bundle_name,
                tags,
                trained_slice_ids,
            ),
            "models/add",
        )
        model_id = response.get("model_id", None)
        if not model_id:
            raise ModelCreationError(response.get("error"))

        return Model(
            model_id=model_id,
            name=name,
            reference_id=reference_id,
            metadata=metadata,
            bundle_name=bundle_name,
            client=self,
            tags=tags,
            trained_slice_ids=trained_slice_ids,
        )

    def create_launch_model(
        self,
        name: str,
        reference_id: str,
        bundle_args: Dict[str, Any],
        metadata: Optional[Dict] = None,
        trained_slice_ids: Optional[List[str]] = None,
    ) -> Model:
        """
        Adds a :class:`Model` to Nucleus, as well as a Launch bundle from a given function.

        Parameters:
            name: A human-readable name for the model.
            reference_id: Unique, user-controlled ID for the model. This can be
              used, for example, to link to an external storage of models which
              may have its own ID scheme.
            bundle_args: Dict of kwargs for creating a Launch bundle. See the
              note below for supported keys.
            metadata: An arbitrary dictionary of additional data about this model
              that can be stored and retrieved. For example, you can store information
              about the hyperparameters used in training this model.

        Returns:
            :class:`Model`: The newly created model as an object.

        .. note::

            A bundle consists of exactly ``{predict_fn_or_cls}``,
            ``{load_predict_fn + model}``, or
            ``{load_predict_fn + load_model_fn}``. The exact keys depend on
            the Launch client version (use ``env_params`` for v0.x, or
            ``pytorch_image_tag``/``tensorflow_version`` otherwise).

            Supported ``bundle_args`` keys:

            - ``model_bundle_name``: Unique identifier for the bundle.
            - ``predict_fn_or_cls``: End-to-end callable for inference.
            - ``model``: Trained neural network, e.g. a PyTorch module.
            - ``load_predict_fn``: Returns an inference function given a model.
            - ``load_model_fn``: Loads a model.
            - ``bundle_url``: Self-hosted mode only. Desired bundle location.
            - ``requirements``: List of pip packages.
            - ``app_config``: YAML dict or local path.
            - ``env_params``: Launch v0 framework/CUDA config.
            - ``globals_copy``: Global symbol table (from ``globals()``).
            - ``pytorch_image_tag``: Launch v1 + PyTorch image tag.
            - ``tensorflow_version``: Launch v1 + TensorFlow version.
        """
        from launch import LaunchClient

        launch_client = LaunchClient(api_key=self.api_key)

        model_exists = any(model.name == name for model in self.list_models())
        bundle_exists = any(
            bundle.name == name + "-nucleus-autogen"
            for bundle in launch_client.list_model_bundles()
        )

        if bundle_exists or model_exists:
            raise ModelCreationError(
                "Bundle with the given name already exists, please try a different name"
            )

        kwargs = {
            "model_bundle_name": name + "-nucleus-autogen",
            **bundle_args,
        }
        if hasattr(launch_client, "create_model_bundle_from_callable_v2"):
            # Launch client is >= 1.0.0
            bundle = launch_client.create_model_bundle_from_callable_v2(
                **kwargs
            )
            bundle_name = (
                bundle.name
            )  # both v0 and v1 have a .name field but are different types
        else:
            bundle = launch_client.create_model_bundle(**kwargs)
            bundle_name = bundle.name
        return self.create_model(
            name,
            reference_id,
            metadata,
            bundle_name,
            trained_slice_ids=trained_slice_ids,
        )

    def create_launch_model_from_dir(
        self,
        name: str,
        reference_id: str,
        bundle_from_dir_args: Dict[str, Any],
        metadata: Optional[Dict] = None,
        trained_slice_ids: Optional[List[str]] = None,
    ) -> Model:
        """Adds a :class:`Model` to Nucleus, as well as a Launch bundle from a directory.

        Parameters:
            name: A human-readable name for the model.
            reference_id: Unique, user-controlled ID for the model. This can be
              used, for example, to link to an external storage of models which
              may have its own id scheme.
            bundle_from_dir_args: Dict of kwargs for creating a bundle from
              local directories. See the note below for supported keys.
            metadata: An arbitrary dictionary of additional data about this model
              that can be stored and retrieved. For example, you can store information
              about the hyperparameters used in training this model.

        Returns:
            :class:`Model`: The newly created model as an object.

        .. note::

            Code from one or more local filesystem folders is packaged into a
            zip and uploaded to Scale Launch. Contents are unzipped relative to
            the server-side ``PYTHONPATH``, so module paths should reflect the
            directory structure (e.g. ``my_module.my_file.f``). The exact keys
            depend on the Launch client version (use ``env_params`` for v0.x,
            or ``pytorch_image_tag``/``tensorflow_version`` otherwise).

            Supported ``bundle_from_dir_args`` keys:

            - ``model_bundle_name``: Unique identifier for the bundle.
            - ``base_paths``: Local dirs containing the bundle code.
            - ``requirements_path``: Path to a ``requirements.txt`` file.
            - ``env_params``: Launch v0 framework/CUDA config.
            - ``load_predict_fn_module_path``: Module path for inference fn.
            - ``load_model_fn_module_path``: Module path for model loader.
            - ``app_config``: YAML dict or local path.
            - ``pytorch_image_tag``: Launch v1 + PyTorch image tag.
            - ``tensorflow_version``: Launch v1 + TensorFlow version.

        .. note::

            For example, given this directory structure::

                my_root/
                    my_module1/
                        __init__.py
                        ...files and directories
                        my_inference_file.py
                    my_module2/
                        __init__.py
                        ...files and directories

            Calling with ``base_paths=["my_module1", "my_module2"]`` creates a
            zip without the root directory. Contents are unzipped relative to
            the server-side ``PYTHONPATH``. If ``my_inference_file.py`` has
            ``def f(...)`` as the inference loading function, then
            ``load_predict_fn_module_path`` should be
            ``my_module1.my_inference_file.f``.
        """
        from launch import LaunchClient

        launch_client = LaunchClient(api_key=self.api_key)

        model_exists = any(model.name == name for model in self.list_models())
        bundle_exists = any(
            bundle.name == name + "-nucleus-autogen"
            for bundle in launch_client.list_model_bundles()
        )

        if bundle_exists or model_exists:
            raise ModelCreationError(
                "Bundle with the given name already exists, please try a different name"
            )

        kwargs = {
            "model_bundle_name": name + "-nucleus-autogen",
            **bundle_from_dir_args,
        }

        if hasattr(launch_client, "create_model_bundle_from_dirs_v2"):
            # Launch client is >= 1.0.0, use new fn
            bundle = launch_client.create_model_bundle_from_dirs_v2(**kwargs)
            # Different code paths give different types for bundle, although both have a .name field
            bundle_name = bundle.name
        else:
            # Launch client is < 1.0.0
            bundle = launch_client.create_model_bundle_from_dirs(**kwargs)
            bundle_name = bundle.name

        return self.create_model(
            name,
            reference_id,
            metadata,
            bundle_name,
            trained_slice_ids=trained_slice_ids,
        )

    @deprecated(
        "Model runs have been deprecated and will be removed. Use a Model instead"
    )
    def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
        response = self.make_request(
            payload, f"dataset/{dataset_id}/modelRun/create"
        )
        if response.get(STATUS_CODE_KEY, None):
            raise ModelRunCreationError(response.get("error"))

        return ModelRun(
            response[MODEL_RUN_ID_KEY], dataset_id=dataset_id, client=self
        )

    @deprecated(
        "Model runs have been deprecated and will be removed. Use a Model instead."
    )
    def commit_model_run(
        self, model_run_id: str, payload: Optional[dict] = None
    ):
        # TODO: deprecate ModelRun. this should be renamed to calculate_evaluation_metrics
        #   or completely removed in favor of Model class methods
        if payload is None:
            payload = {}
        return self.make_request(payload, f"modelRun/{model_run_id}/commit")

    @deprecated(msg="Prefer calling Dataset.info() directly.")
    def dataset_info(self, dataset_id: str):
        dataset = self.get_dataset(dataset_id)
        return dataset.info()

    @deprecated(
        "Model runs have been deprecated and will be removed. Use a Model instead."
    )
    def model_run_info(self, model_run_id: str):
        # TODO: deprecate ModelRun
        return self.make_request(
            {}, f"modelRun/{model_run_id}/info", requests.get
        )

    @deprecated("Prefer calling Dataset.refloc instead.")
    @sanitize_string_args
    def dataitem_ref_id(self, dataset_id: str, reference_id: str):
        # TODO: deprecate in favor of Dataset.refloc invocation
        return self.make_request(
            {}, f"dataset/{dataset_id}/refloc/{reference_id}", requests.get
        )

    @deprecated("Prefer calling Dataset.predictions_refloc instead.")
    @sanitize_string_args
    def predictions_ref_id(
        self, model_run_id: str, ref_id: str, dataset_id: Optional[str] = None
    ):
        if dataset_id:
            raise RuntimeError(
                "Need to pass a dataset id. Or use Dataset.predictions_refloc."
            )
        # TODO: deprecate ModelRun
        m_run = self.get_model_run(model_run_id, dataset_id)
        return m_run.refloc(ref_id)

    @deprecated("Prefer calling Dataset.iloc instead.")
    def dataitem_iloc(self, dataset_id: str, i: int):
        # TODO: deprecate in favor of Dataset.iloc invocation
        return self.make_request(
            {}, f"dataset/{dataset_id}/iloc/{i}", requests.get
        )

    @deprecated("Prefer calling Dataset.predictions_iloc instead.")
    def predictions_iloc(self, model_run_id: str, i: int):
        # TODO: deprecate ModelRun
        return self.make_request(
            {}, f"modelRun/{model_run_id}/iloc/{i}", requests.get
        )

    @deprecated("Prefer calling Dataset.loc instead.")
    def dataitem_loc(self, dataset_id: str, dataset_item_id: str):
        # TODO: deprecate in favor of Dataset.loc invocation
        return self.make_request(
            {}, f"dataset/{dataset_id}/loc/{dataset_item_id}", requests.get
        )

    @deprecated("Prefer calling Dataset.predictions_loc instead.")
    def predictions_loc(self, model_run_id: str, dataset_item_id: str):
        # TODO: deprecate ModelRun
        return self.make_request(
            {}, f"modelRun/{model_run_id}/loc/{dataset_item_id}", requests.get
        )

    @deprecated("Prefer calling Dataset.create_slice instead.")
    def create_slice(self, dataset_id: str, payload: dict) -> Slice:
        # TODO: deprecate in favor of Dataset.create_slice
        dataset = self.get_dataset(dataset_id)
        return dataset.create_slice(payload["name"], payload["reference_ids"])

    def get_slice(self, slice_id: str) -> Slice:
        # TODO: migrate to Dataset method and deprecate
        """Returns a slice object by Nucleus-generated ID.

        Parameters:
            slice_id: Nucleus-generated slice ID (starts with ``slc_``). This can
              be retrieved via :meth:`Dataset.slices` or a Nucleus dashboard URL.

        Returns:
            :class:`Slice`: The Nucleus slice as an object.
        """
        return Slice(slice_id, self)

    @deprecated("Prefer calling Slice.info instead.")
    def slice_info(self, slice_id: str) -> dict:
        # TODO: deprecate in favor of Slice.info
        response = self.make_request(
            {},
            f"slice/{slice_id}",
            requests_command=requests.get,
        )
        return response

    def delete_slice(self, slice_id: str) -> dict:
        # TODO: migrate to Dataset method and deprecate
        """Deletes slice from Nucleus.

        Parameters:
            slice_id: Nucleus-generated slice ID (starts with ``slc_``). This can
              be retrieved via :meth:`Dataset.slices` or a Nucleus dashboard URL.

        Returns:
            Empty payload response.
        """
        response = self.make_request(
            {},
            f"slice/{slice_id}",
            requests_command=requests.delete,
        )
        return response

    @deprecated("Prefer calling Dataset.delete_annotations instead.")
    def delete_annotations(
        self,
        dataset_id: str,
        reference_ids: Optional[list] = None,
        keep_history=True,
    ) -> AsyncJob:
        dataset = self.get_dataset(dataset_id)
        return dataset.delete_annotations(reference_ids, keep_history)

    def append_to_slice(
        self,
        slice_id: str,
        reference_ids: List[str],
        dataset_id: str,
    ) -> dict:
        # TODO: migrate to Slice method and deprecate
        """Appends dataset items or scenes to an existing slice.

        Parameters:
            slice_id: Nucleus-generated dataset ID (starts with ``slc_``). This can
              be retrieved via :meth:`Dataset.slices` or a Nucleus dashboard URL.
            reference_ids: List of user-defined reference IDs of dataset items or scenes
              to append to the slice.
            dataset_id: ID of dataset this slice belongs to.

        Returns:
            Empty payload response.
        """

        response = self.make_request(
            {REFERENCE_IDS_KEY: reference_ids, DATASET_ID_KEY: dataset_id},
            f"slice/{slice_id}/append",
        )
        return response

    def list_autotags(self, dataset_id: str) -> List[dict]:
        # TODO: deprecate in favor of Dataset.list_autotags invocation
        response = self.make_request(
            {},
            f"{dataset_id}/list_autotags",
            requests_command=requests.get,
        )
        return response[AUTOTAGS_KEY] if AUTOTAGS_KEY in response else response

    def delete_autotag(self, autotag_id: str) -> dict:
        # TODO: migrate to Dataset method (use autotag name, not id) and deprecate
        """Deletes an autotag by ID.

        Parameters:
            autotag_id: Nucleus-generated autotag ID (starts with ``tag_``). This can
              be retrieved via :meth:`list_autotags` or a Nucleus dashboard URL.

        Returns:
            Empty payload response.
        """
        return self.make_request({}, f"autotag/{autotag_id}", requests.delete)

    def get_autotag_refinement_metrics(self, autotag_id: str) -> dict:
        """Retrieves refinement metrics for an autotag by ID.

        Parameters:
            autotag_id: Nucleus-generated autotag ID (starts with ``tag_``). This can
              be retrieved via :meth:`list_autotags` or a Nucleus dashboard URL.

        Returns:
            Response payload::

                {
                    "total_refinement_steps": int
                    "average_positives_selected_per_refinement": int
                    "average_ms_taken_in_refinement": float
                }
        """
        return self.make_request(
            {}, f"autotag/{autotag_id}/refinementMetrics", requests.get
        )

    def delete_model(self, model_id: str) -> dict:
        """Deletes a model by ID.

        Parameters:
            model_id: Nucleus-generated model ID (starts with ``prj_``). This can
              be retrieved via :meth:`list_models` or a Nucleus dashboard URL.

        Returns:
            Empty payload response.
        """
        response = self.make_request(
            {},
            f"model/{model_id}",
            requests_command=requests.delete,
        )
        return response

    def download_pointcloud_task(
        self, task_id: str, frame_num: int
    ) -> List[Union[Point3D, LidarPoint]]:
        """
        Download the lidar point cloud data for a given task and frame number.

        Parameters:
            task_id: download point cloud for this particular task
            frame_num: download point cloud for this particular frame

        Returns:
            List of Point3D objects

        """

        response = self.make_request(
            payload={},
            route=f"task/{task_id}/frame/{frame_num}",
            requests_command=requests.get,
        )
        points = response.get(POINTS_KEY, None)
        if points is None or len(points) == 0:
            raise Exception("Response has invalid payload")

        sample_point = points[0]
        if I_KEY in sample_point.keys():
            return [LidarPoint.from_json(pt) for pt in points]

        return [Point3D.from_json(pt) for pt in points]

    def download_pointcloud_tasks(
        self, task_ids: List[str], frame_num: int
    ) -> Dict[str, List[Union[Point3D, LidarPoint]]]:
        """
        Download the lidar point cloud data for a given set of tasks and frame number.

        Parameters:
            task_ids: list of task ids to fetch data from
            frame_num: download point cloud for this particular frame

        Returns:
            A dictionary from task_id to list of Point3D objects

        """
        endpoints = [
            f"task/{task_id}/frame/{frame_num}" for task_id in task_ids
        ]
        progressbar = self.tqdm_bar(
            total=len(endpoints),
            desc="Downloading pointcloud tasks",
        )
        results = make_multiple_requests_concurrently(
            client=self,
            requests=endpoints,
            route=None,
            progressbar=progressbar,
        )
        resp = {}

        for result in results:
            req, data = result
            task_id = req.split("/")[1]  # task/<task id>/frame/1 => task_id
            points = data.get(POINTS_KEY, None)
            if points is None or len(points) == 0:
                raise Exception("Response has invalid payload")

            sample_point = points[0]
            if I_KEY in sample_point.keys():
                resp[task_id] = [LidarPoint.from_json(pt) for pt in points]
            else:
                resp[task_id] = [Point3D.from_json(pt) for pt in points]

        return resp

    @deprecated("Prefer calling Dataset.create_custom_index instead.")
    def create_custom_index(
        self, dataset_id: str, embeddings_urls: list, embedding_dim: int
    ):
        # TODO: deprecate in favor of Dataset.create_custom_index invocation
        dataset = self.get_dataset(dataset_id)
        return dataset.create_custom_index(
            embeddings_urls=embeddings_urls, embedding_dim=embedding_dim
        )

    @deprecated("Prefer calling Dataset.delete_custom_index instead.")
    def delete_custom_index(self, dataset_id: str, image: bool):
        # TODO: deprecate in favor of Dataset.delete_custom_index invocation
        return self.make_request(
            {"image": image},
            f"indexing/{dataset_id}",
            requests_command=requests.delete,
        )

    @deprecated("Prefer calling Dataset.set_primary_index instead.")
    def set_primary_index(self, dataset_id: str, image: bool, custom: bool):
        # TODO: deprecate in favor of Dataset.set_primary_index invocation
        return self.make_request(
            {"image": image, "custom": custom},
            f"indexing/{dataset_id}/setPrimary",
            requests_command=requests.post,
        )

    @deprecated("Prefer calling Dataset.set_continuous_indexing instead.")
    def set_continuous_indexing(self, dataset_id: str, enable: bool = True):
        # TODO: deprecate in favor of Dataset.set_continuous_indexing invocation
        return self.make_request(
            {INDEX_CONTINUOUS_ENABLE_KEY: enable},
            f"indexing/{dataset_id}/setContinuous",
            requests_command=requests.post,
        )

    @deprecated("Prefer calling Dataset.create_image_index instead.")
    def create_image_index(self, dataset_id: str):
        # TODO: deprecate in favor of Dataset.create_image_index invocation
        return self.make_request(
            {},
            f"indexing/{dataset_id}/internal/image",
            requests_command=requests.post,
        )

    @deprecated("Prefer calling Dataset.create_object_index instead.")
    def create_object_index(
        self, dataset_id: str, model_run_id: str, gt_only: bool
    ):
        # TODO: deprecate in favor of Dataset.create_object_index invocation
        payload: Dict[str, Union[str, bool]] = {}
        if model_run_id:
            payload["model_run_id"] = model_run_id
        elif gt_only:
            payload["ingest_gt_only"] = True
        return self.make_request(
            payload,
            f"indexing/{dataset_id}/internal/object",
            requests_command=requests.post,
        )

    def delete(self, route: str):
        return self.connection.delete(route)

    def get(self, route: str):
        return self.connection.get(route)

    def post(self, payload: dict, route: str):
        return self.connection.post(payload, route)

    def put(self, payload: dict, route: str):
        return self.connection.put(payload, route)

    # TODO: Fix return type, can be a list as well. Brings on a lot of mypy errors ...
    def make_request(
        self,
        payload: Optional[dict],
        route: str,
        requests_command=requests.post,
        return_raw_response: bool = False,
    ) -> Union[dict, Any]:
        """Makes a request to a Nucleus API endpoint.

        Logs a warning if not successful.

        Parameters:
            payload: Given request payload.
            route: Route for the request.
            requests_command: ``requests.post``, ``requests.get``, or ``requests.delete``.
            return_raw_response: Whether to return the raw response object.

        Returns:
            Response payload as JSON dict or request object.
        """
        if payload is None:
            payload = {}
        if requests_command is requests.get:
            if payload:
                print(
                    "Received defined payload with GET request! Will ignore payload"
                )
            payload = None
        return self.connection.make_request(payload, route, requests_command, return_raw_response)  # type: ignore

    def _set_api_key(self, api_key):
        """Fetch API key from environment variable NUCLEUS_API_KEY if not set"""
        api_key = (
            api_key if api_key else os.environ.get("NUCLEUS_API_KEY")
        )
        if api_key is None:
            raise NoAPIKey()

        return api_key

    @staticmethod
    def valid_dirname(dirname) -> str:
        """Validates that a directory exists.

        Parameters:
            dirname: Path of directory.

        Returns:
            Existing directory path.
        """
        # ensures path ends with a slash
        _dirname = os.path.join(os.path.expanduser(dirname), "")
        if not os.path.exists(_dirname):
            raise ValueError(
                f"Given directory name: {dirname} does not exists. Searched in {_dirname}"
            )
        return _dirname

    def create_dataset_from_dir(
        self,
        dirname: str,
        dataset_name: Optional[str] = None,
        use_privacy_mode: bool = False,
        privacy_mode_proxy: str = "",
        allowed_file_types: Tuple[str, ...] = ("png", "jpg", "jpeg"),
        skip_size_warning: bool = False,
    ) -> Dataset:
        """
        Create a dataset by recursively crawling through a directory.
        A DatasetItem will be created for each unique image found.

        Parameters:
            dirname: Where to look for image files, recursively
            dataset_name: If none is given, the parent folder name is used
            use_privacy_mode: Whether the dataset should be treated as privacy
            privacy_mode_proxy: Endpoint that serves image files for privacy mode, ignore if not using privacy mode.
                The proxy should work based on the relative path of the images in the directory.
            allowed_file_types: Which file type extensions to search for, ie: ('jpg', 'png')
            skip_size_warning: If False, it will throw an error if the script globs more than 500 images. This is a safety check in case the dirname has a typo, and grabs too much data.
        """
        existing_dirname = self.valid_dirname(dirname)
        folder_name = os.path.basename(existing_dirname.rstrip("/"))
        dataset_name = dataset_name or folder_name
        dataset = self.create_dataset(
            name=dataset_name, use_privacy_mode=use_privacy_mode
        )
        dataset.add_items_from_dir(
            existing_dirname=existing_dirname,
            privacy_mode_proxy=privacy_mode_proxy,
            allowed_file_types=allowed_file_types,
            skip_size_warning=skip_size_warning,
        )
        return dataset