95 lines
2.6 KiB
Python
95 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
class FakeTensor:
|
|
def __init__(self, array):
|
|
self.array = np.asarray(array, dtype=np.float32)
|
|
|
|
def detach(self):
|
|
return self
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def float(self):
|
|
return self
|
|
|
|
def numpy(self):
|
|
return self.array
|
|
|
|
|
|
class FakeBatch:
|
|
def __init__(self, arrays):
|
|
self.items = [FakeTensor(array) for array in arrays]
|
|
|
|
def __len__(self):
|
|
return len(self.items)
|
|
|
|
def __getitem__(self, index):
|
|
return self.items[index]
|
|
|
|
|
|
@pytest.fixture()
|
|
def depth_module(tmp_path, monkeypatch):
|
|
folder_paths = types.ModuleType("folder_paths")
|
|
folder_paths.get_temp_directory = lambda: str(tmp_path)
|
|
|
|
def get_save_image_path(prefix, output_dir, width, height):
|
|
assert width > 0 and height > 0
|
|
return output_dir, f"{prefix}_%batch_num%", 1, "", prefix
|
|
|
|
folder_paths.get_save_image_path = get_save_image_path
|
|
monkeypatch.setitem(sys.modules, "folder_paths", folder_paths)
|
|
|
|
module_path = Path(__file__).parents[1] / "__init__.py"
|
|
spec = importlib.util.spec_from_file_location("depth_viewer_test_module", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader
|
|
spec.loader.exec_module(module)
|
|
return module, tmp_path
|
|
|
|
|
|
def test_processes_full_batch_and_broadcasts_single_depth(depth_module):
|
|
module, output_dir = depth_module
|
|
reference = FakeBatch(
|
|
[
|
|
np.zeros((8, 12, 3)),
|
|
np.ones((8, 12, 3)),
|
|
]
|
|
)
|
|
depth = FakeBatch([np.full((8, 12, 1), 0.5)])
|
|
|
|
result = module.DepthViewer().process_images(reference, depth)["ui"]
|
|
|
|
assert len(result["reference_image"]) == 2
|
|
assert len(result["depth_map"]) == 2
|
|
for descriptor in result["reference_image"] + result["depth_map"]:
|
|
assert descriptor["type"] == "temp"
|
|
assert (output_dir / descriptor["filename"]).is_file()
|
|
|
|
|
|
def test_rejects_incompatible_batch_sizes(depth_module):
|
|
module, _ = depth_module
|
|
reference = FakeBatch([np.zeros((4, 4, 3)) for _ in range(2)])
|
|
depth = FakeBatch([np.zeros((4, 4, 1)) for _ in range(3)])
|
|
|
|
with pytest.raises(ValueError, match="Batch sizes must match"):
|
|
module.DepthViewer().process_images(reference, depth)
|
|
|
|
|
|
def test_sanitizes_non_finite_pixels(depth_module):
|
|
module, _ = depth_module
|
|
image = FakeTensor([[[np.nan, np.inf, -np.inf]]])
|
|
|
|
converted = np.asarray(module._as_pil(image))
|
|
|
|
assert converted.tolist() == [[[0, 255, 0]]]
|