diff --git a/.gitignore b/.gitignore index 28afce358a..25f2ceabcf 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ develop-eggs/ dist/ downloads/ !nerfstudio/scripts/downloads/ +!tests/scripts/downloads/ eggs/ .eggs/ lib/ diff --git a/nerfstudio/scripts/downloads/download_data.py b/nerfstudio/scripts/downloads/download_data.py index a2c9d27584..cf5a2fccd5 100644 --- a/nerfstudio/scripts/downloads/download_data.py +++ b/nerfstudio/scripts/downloads/download_data.py @@ -584,6 +584,32 @@ def main(self) -> None: ... ] +# Add mipnerf360 (Amazon S3 / S3-compatible) subcommand if boto3 is installed. +try: + import boto3 as _boto3 # noqa: F401 +except ImportError: + _boto3 = None + +if _boto3 is not None: + from nerfstudio.scripts.downloads.mipnerf360_download import Mipnerf360Download + + Commands = Union[ + Commands, + Annotated[Mipnerf360Download, tyro.conf.subcommand(name="mipnerf360")], + ] +else: + Commands = Union[ + Commands, + Annotated[ + NotInstalled, + tyro.conf.subcommand( + name="mipnerf360", + description="**Not installed.** Downloading the Mip-NeRF 360 dataset requires `pip install boto3`.", + ), + ], + ] + + def main( dataset: DatasetDownload, ): diff --git a/nerfstudio/scripts/downloads/mipnerf360_download.py b/nerfstudio/scripts/downloads/mipnerf360_download.py new file mode 100644 index 0000000000..3ced6cfe60 --- /dev/null +++ b/nerfstudio/scripts/downloads/mipnerf360_download.py @@ -0,0 +1,168 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Mip-NeRF 360 dataset downloader. + +Downloads the Mip-NeRF 360 scenes from an anonymously-readable bucket using the +Amazon S3 API. The endpoint and bucket are plumbing: point +``AWS_S3_ENDPOINT_URL`` (or ``--bucket``) at Amazon S3 or any S3-compatible +object store (for example Backblaze B2, Cloudflare R2, or MinIO) and the same +code path works. + +Registered as the ``mipnerf360`` subcommand under ``ns-download-data``. +""" + +import os +import sys +from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path + +from nerfstudio.scripts.downloads.utils import DatasetDownload +from nerfstudio.utils.rich_utils import CONSOLE + +try: + import boto3 + from botocore import UNSIGNED + from botocore.client import Config +except ImportError: + print("boto3 is required for the Mip-NeRF 360 downloader. Install it with `pip install boto3`.") + sys.exit(1) + + +# Placeholder S3 endpoint. Override with the standard AWS_S3_ENDPOINT_URL env +# var to target Amazon S3 or any S3-compatible host. +DEFAULT_S3_ENDPOINT_URL = "https://your-s3-endpoint.example.com" + +# Placeholder bucket name. Replace with the real public bucket once provisioned. +DEFAULT_BUCKET = "nerfstudio-mipnerf360" + +# Mip-NeRF 360 scene names, matching the canonical Barron et al. release. +MIPNERF360_SCENES = ( + "bicycle", + "bonsai", + "counter", + "garden", + "kitchen", + "room", + "stump", + "flowers", + "treehill", +) + + +def _nerfstudio_user_agent_extra() -> str: + """Return the ``nerfstudio/`` user agent fragment. + + Falls back to ``nerfstudio/dev`` when the package metadata is unavailable, + for example when running from a source checkout that was never installed. + """ + try: + return f"nerfstudio/{version('nerfstudio')}" + except PackageNotFoundError: + return "nerfstudio/dev" + + +def _resolve_endpoint_url() -> str: + """Pick the S3 endpoint URL. + + Precedence: ``AWS_S3_ENDPOINT_URL`` env var, otherwise the placeholder + default. + """ + return os.environ.get("AWS_S3_ENDPOINT_URL", DEFAULT_S3_ENDPOINT_URL) + + +def _build_s3_client(anonymous: bool = True): + """Build a boto3 S3 client for the configured S3-compatible endpoint. + + Args: + anonymous: If True (the default for public dataset downloads), uses + unsigned requests so users do not need credentials. + + Returns: + A configured ``boto3.client('s3')`` instance. + """ + config = Config( + user_agent_extra=_nerfstudio_user_agent_extra(), + # Virtual-hosted-style addressing works across Amazon S3 and S3-compatible stores. + s3={"addressing_style": "virtual"}, + signature_version=UNSIGNED if anonymous else None, + ) + return boto3.client("s3", endpoint_url=_resolve_endpoint_url(), config=config) + + +def _download_prefix(client, bucket: str, prefix: str, output_dir: Path) -> int: + """Download every object under ``prefix`` from ``bucket`` into ``output_dir``. + + Returns the number of files downloaded. + """ + paginator = client.get_paginator("list_objects_v2") + count = 0 + for page in paginator.paginate(Bucket=bucket, Prefix=prefix): + for obj in page.get("Contents", []) or []: + key = obj["Key"] + # Strip the prefix so the on-disk layout mirrors the scene contents, + # not the bucket nesting. + rel = key[len(prefix) :].lstrip("/") + if not rel: + continue + dest = output_dir / rel + dest.parent.mkdir(parents=True, exist_ok=True) + client.download_file(bucket, key, str(dest)) + count += 1 + return count + + +@dataclass +class Mipnerf360Download(DatasetDownload): + """Download the Mip-NeRF 360 dataset over the Amazon S3 API. + + Uses anonymous (unsigned) reads by default. Honors ``AWS_S3_ENDPOINT_URL`` + if set, so the same command works against Amazon S3 or any S3-compatible + object store (for example Backblaze B2, Cloudflare R2, or MinIO). + """ + + capture_name: str = "all" + """Which Mip-NeRF 360 scene to fetch ("all" or one of the scene names).""" + + bucket: str = DEFAULT_BUCKET + """Name of the bucket to download from. Override to use a different mirror.""" + + anonymous: bool = True + """Use unsigned requests (public buckets). Set to False to use S3 credentials from your env.""" + + def download(self, save_dir: Path) -> None: + """Download the requested scenes into ``save_dir / mipnerf360 / ``.""" + if self.capture_name == "all": + scenes = list(MIPNERF360_SCENES) + elif self.capture_name in MIPNERF360_SCENES: + scenes = [self.capture_name] + else: + CONSOLE.print( + f"[bold red]Unknown capture '{self.capture_name}'. Valid choices: {list(MIPNERF360_SCENES) + ['all']}." + ) + sys.exit(1) + + endpoint = _resolve_endpoint_url() + CONSOLE.print(f"Downloading from endpoint [bold]{endpoint}[/bold], bucket [bold]{self.bucket}[/bold].") + client = _build_s3_client(anonymous=self.anonymous) + + for i, scene in enumerate(scenes): + prefix = f"{scene}/" + output_dir = save_dir / "mipnerf360" / scene + output_dir.mkdir(parents=True, exist_ok=True) + CONSOLE.print(f"[mipnerf360 {i + 1: >2d}/{len(scenes)}] downloading scene '{scene}' to {output_dir}") + n = _download_prefix(client, self.bucket, prefix, output_dir) + CONSOLE.print(f" done. {n} file(s) downloaded.") diff --git a/tests/scripts/downloads/__init__.py b/tests/scripts/downloads/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/scripts/downloads/test_mipnerf360_download.py b/tests/scripts/downloads/test_mipnerf360_download.py new file mode 100644 index 0000000000..d20791a30b --- /dev/null +++ b/tests/scripts/downloads/test_mipnerf360_download.py @@ -0,0 +1,123 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Mip-NeRF 360 (Amazon S3 API) dataset downloader.""" + +from unittest import mock + +import pytest + + +def test_resolve_endpoint_url_default(monkeypatch): + """Default endpoint is the placeholder S3 URL.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + monkeypatch.delenv("AWS_S3_ENDPOINT_URL", raising=False) + assert mipnerf360_download._resolve_endpoint_url() == mipnerf360_download.DEFAULT_S3_ENDPOINT_URL + + +def test_resolve_endpoint_url_env_override(monkeypatch): + """AWS_S3_ENDPOINT_URL takes precedence over the default.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + monkeypatch.setenv("AWS_S3_ENDPOINT_URL", "https://s3.example.test") + assert mipnerf360_download._resolve_endpoint_url() == "https://s3.example.test" + + +def test_build_s3_client_sets_user_agent_and_addressing(monkeypatch): + """boto3.client is called with the nerfstudio user agent extra and virtual addressing.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + fake_client = mock.MagicMock() + with mock.patch.object(mipnerf360_download, "boto3") as fake_boto3: + fake_boto3.client.return_value = fake_client + result = mipnerf360_download._build_s3_client(anonymous=True) + + assert result is fake_client + fake_boto3.client.assert_called_once() + args, kwargs = fake_boto3.client.call_args + assert args[0] == "s3" + assert "endpoint_url" in kwargs + config = kwargs["config"] + assert config.user_agent_extra.startswith("nerfstudio/") + assert config.s3 == {"addressing_style": "virtual"} + + +def test_user_agent_extra_uses_package_version(monkeypatch): + """The user agent fragment is nerfstudio/ when metadata is available.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + monkeypatch.setattr(mipnerf360_download, "version", lambda name: "9.9.9") + assert mipnerf360_download._nerfstudio_user_agent_extra() == "nerfstudio/9.9.9" + + +def test_user_agent_extra_falls_back_to_dev(monkeypatch): + """The user agent fragment falls back to nerfstudio/dev without package metadata.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + def _raise(name): + raise mipnerf360_download.PackageNotFoundError(name) + + monkeypatch.setattr(mipnerf360_download, "version", _raise) + assert mipnerf360_download._nerfstudio_user_agent_extra() == "nerfstudio/dev" + + +def test_build_s3_client_anonymous_uses_unsigned(): + """anonymous=True selects the UNSIGNED signature version.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + with mock.patch.object(mipnerf360_download, "boto3") as fake_boto3: + mipnerf360_download._build_s3_client(anonymous=True) + config = fake_boto3.client.call_args.kwargs["config"] + assert config.signature_version == mipnerf360_download.UNSIGNED + + +def test_build_s3_client_signed_when_not_anonymous(): + """anonymous=False does not force the UNSIGNED signature version.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + with mock.patch.object(mipnerf360_download, "boto3") as fake_boto3: + mipnerf360_download._build_s3_client(anonymous=False) + config = fake_boto3.client.call_args.kwargs["config"] + assert config.signature_version is None + + +@pytest.mark.parametrize("scene", ["bicycle", "garden", "treehill"]) +def test_mipnerf360_download_single_scene_invokes_client(monkeypatch, tmp_path, scene): + """Downloading a specific scene paginates and downloads each listed object.""" + from nerfstudio.scripts.downloads import mipnerf360_download + + fake_client = mock.MagicMock() + paginator = mock.MagicMock() + paginator.paginate.return_value = [ + {"Contents": [{"Key": f"{scene}/images/0001.png"}, {"Key": f"{scene}/transforms.json"}]}, + ] + fake_client.get_paginator.return_value = paginator + + monkeypatch.setattr(mipnerf360_download, "_build_s3_client", lambda anonymous=True: fake_client) + + dl = mipnerf360_download.Mipnerf360Download(capture_name=scene, bucket="nerfstudio-mipnerf360") + dl.download(tmp_path) + + paginator.paginate.assert_called_once_with(Bucket="nerfstudio-mipnerf360", Prefix=f"{scene}/") + assert fake_client.download_file.call_count == 2 + + +def test_mipnerf360_download_unknown_scene_exits(monkeypatch, tmp_path): + """An unknown capture_name calls sys.exit(1).""" + from nerfstudio.scripts.downloads import mipnerf360_download + + dl = mipnerf360_download.Mipnerf360Download(capture_name="not-a-scene") + with pytest.raises(SystemExit): + dl.download(tmp_path)