Spaces:
Running on Zero
Running on Zero
File size: 2,653 Bytes
41ff959 a6825eb 41ff959 a6825eb 41ff959 a6825eb 41ff959 a6825eb 41ff959 | 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 | from pathlib import Path
import cv2
import h5py
import numpy as np
import pytest
import torch
from src.demo.infer_single_image import _read_depth_array, load_depth
def test_prompt_depth_loads_plain_npy_and_npz(tmp_path: Path) -> None:
expected = np.arange(12, dtype=np.float32).reshape(3, 4)
npy_path = tmp_path / "depth.npy"
npz_path = tmp_path / "depth.npz"
np.save(npy_path, expected)
np.savez(npz_path, depth=expected)
np.testing.assert_array_equal(_read_depth_array(npy_path), expected)
np.testing.assert_array_equal(_read_depth_array(npz_path), expected)
def test_prompt_depth_loads_sparse_npz(tmp_path: Path) -> None:
depth_path = tmp_path / "depth.npz"
mask = np.array([[False, True], [True, False]])
np.savez(depth_path, mask=mask, value=np.array([2.0, 4.0], dtype=np.float32))
actual = _read_depth_array(depth_path)
np.testing.assert_array_equal(
actual,
np.array([[0.0, 2.0], [4.0, 0.0]], dtype=np.float32),
)
def test_prompt_depth_loads_channel_first_npy(tmp_path: Path) -> None:
depth_path = tmp_path / "depth.npy"
expected = np.arange(12, dtype=np.float32).reshape(3, 4)
np.save(depth_path, expected[None])
np.testing.assert_array_equal(_read_depth_array(depth_path), expected)
def test_prompt_depth_rejects_png(tmp_path: Path) -> None:
depth_path = tmp_path / "depth.png"
depth_path.touch()
with pytest.raises(ValueError, match="Unsupported prompt depth extension"):
_read_depth_array(depth_path)
def test_prompt_depth_loads_nested_hdf5(tmp_path: Path) -> None:
depth_path = tmp_path / "depth.h5"
expected = np.arange(12, dtype=np.float32).reshape(3, 4)
with h5py.File(depth_path, "w") as h5_file:
h5_file.create_dataset("nested/depth", data=expected)
np.testing.assert_array_equal(_read_depth_array(depth_path), expected)
def test_prompt_depth_loads_exr(tmp_path: Path) -> None:
depth_path = tmp_path / "depth.exr"
expected = np.arange(12, dtype=np.float32).reshape(3, 4)
assert cv2.imwrite(str(depth_path), expected)
np.testing.assert_array_equal(_read_depth_array(depth_path), expected)
def test_load_depth_preserves_default_sparse_sampling(tmp_path: Path) -> None:
depth_path = tmp_path / "depth.npy"
np.save(depth_path, np.full((40, 50), 2.0, dtype=np.float32))
np.random.seed(0)
dense_depth, sampled_depth, depth_mask = load_depth(
depth_path=depth_path,
tar_size=(40, 50),
)
assert dense_depth.shape == (1, 1, 40, 50)
assert depth_mask.sum().item() == 2000
assert torch.count_nonzero(sampled_depth).item() == 1500
|