Spaces:
Running on Zero
Running on Zero
File size: 4,906 Bytes
41ff959 a6825eb 41ff959 a6825eb 41ff959 a6825eb 41ff959 a6825eb 41ff959 a6825eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | from argparse import Namespace
from pathlib import Path
import pytest
import src.demo.infer_batch_images as batch_module
from src.demo.infer_batch_images import (
_collect_images,
_resolve_output_dir,
_resolve_prompt_depth_path,
run_batch,
)
def _args(input_path, **overrides) -> Namespace:
values = {
"input_path": input_path,
"mode": "rgb",
"limit": 0,
"prompt_depth": None,
"prompt_depth_dir": None,
}
values.update(overrides)
return Namespace(**values)
def test_collect_images_accepts_single_image(tmp_path) -> None:
image = tmp_path / "input.jpg"
image.touch()
assert _collect_images(_args(image)) == [image]
def test_collect_images_accepts_directory(tmp_path) -> None:
first = tmp_path / "a.png"
second = tmp_path / "b.jpg"
ignored = tmp_path / "notes.txt"
first.touch()
second.touch()
ignored.touch()
assert _collect_images(_args(tmp_path)) == [first, second]
def test_collect_images_does_not_scan_subdirectories(tmp_path) -> None:
image = tmp_path / "input.jpg"
nested = tmp_path / "nested"
nested.mkdir()
nested_image = nested / "nested.jpg"
image.touch()
nested_image.touch()
assert _collect_images(_args(tmp_path)) == [image]
def test_collect_images_rejects_duplicate_stems(tmp_path) -> None:
(tmp_path / "frame.jpg").touch()
(tmp_path / "frame.png").touch()
with pytest.raises(ValueError, match="unique filename stems"):
_collect_images(_args(tmp_path))
def test_lidar_directory_keeps_png_rgb_with_npz_depth(tmp_path) -> None:
image = tmp_path / "frame.png"
depth = tmp_path / "frame.npz"
image.touch()
depth.touch()
assert _collect_images(_args(tmp_path, mode="lidar")) == [image]
def test_default_output_layout() -> None:
rgb_output = _resolve_output_dir(Namespace(output_dir=None, mode="rgb"))
lidar_output = _resolve_output_dir(Namespace(output_dir=None, mode="lidar"))
assert rgb_output == Path("outputs/demo/rgb")
assert lidar_output == Path("outputs/demo/lidar")
def test_lidar_mode_matches_adjacent_depth_by_stem(tmp_path) -> None:
image = tmp_path / "frame.jpg"
depth = tmp_path / "frame.npy"
image.touch()
depth.touch()
args = _args(
image,
mode="lidar",
prompt_depth=None,
prompt_depth_dir=None,
)
assert _resolve_prompt_depth_path(args, image) == depth
def test_lidar_mode_does_not_use_input_png_as_depth(tmp_path) -> None:
image = tmp_path / "frame.png"
image.touch()
args = _args(
image,
mode="lidar",
prompt_depth=None,
prompt_depth_dir=None,
)
with pytest.raises(FileNotFoundError):
_resolve_prompt_depth_path(args, image)
def _run_args(input_path: Path, output_dir: Path, **overrides) -> Namespace:
values = {
"input_path": input_path,
"mode": "rgb",
"checkpoint": None,
"output_dir": output_dir,
"limit": 0,
"overwrite": False,
"device": "auto",
"intrinsics_file": None,
"focal_px": None,
"focal_mm": None,
"prompt_depth": None,
"prompt_depth_dir": None,
"disable_floater_filter": False,
"no_video": True,
"export_html": False,
}
values.update(overrides)
return Namespace(**values)
def test_completed_batch_does_not_load_model(tmp_path, monkeypatch) -> None:
image = tmp_path / "frame.jpg"
output_dir = tmp_path / "outputs"
scene_ply = output_dir / "frame" / "frame.ply"
image.touch()
scene_ply.parent.mkdir(parents=True)
scene_ply.touch()
monkeypatch.setattr(
batch_module,
"load_demo_config",
lambda *_args, **_kwargs: pytest.fail("completed batch loaded the model config"),
)
result = run_batch(_run_args(image, output_dir))
assert result["succeeded"] == 0
assert result["skipped"] == 1
def test_html_only_batch_does_not_load_model(tmp_path, monkeypatch) -> None:
image = tmp_path / "frame.jpg"
output_dir = tmp_path / "outputs"
scene_ply = output_dir / "frame" / "frame.ply"
html = output_dir / "frame" / "frame.html"
image.touch()
scene_ply.parent.mkdir(parents=True)
scene_ply.touch()
monkeypatch.setattr(batch_module.shutil, "which", lambda _command: "/usr/bin/splat-transform")
monkeypatch.setattr(
batch_module,
"load_demo_config",
lambda *_args, **_kwargs: pytest.fail("HTML-only batch loaded the model config"),
)
monkeypatch.setattr(
batch_module,
"_convert_scene_if_requested",
lambda _export_html, _paths, _viewer_settings: html.touch(),
)
result = run_batch(_run_args(image, output_dir, export_html=True))
assert result["succeeded"] == 1
assert result["skipped"] == 0
assert html.exists()
|