Add ns-download-data mipnerf360 subcommand (S3-compatible)

.gitignore

@@ -13,6 +13,7 @@ develop-eggs/
 dist/
 downloads/
 !nerfstudio/scripts/downloads/
+!tests/scripts/downloads/
 eggs/
 .eggs/
 lib/

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,
 ):

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/<version>`` 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 / <scene>``."""
+        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.")

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/<version> 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)