nan
7cd5681a8b
- Implement test to ensure foreground preserves source RGB and Ben2 alpha in WorkerTest. - Remove outdated tests related to background distance and original background depth. - Update geometry node asset to reflect new naming conventions and structure. - Introduce validation for depth plane asset, ensuring correct input/output configurations and attribute handling. - Adjust confidence and edge parameters in geometry node asset to support a wider range of values.
1438 lines
49 KiB
Python
1438 lines
49 KiB
Python
import importlib.util
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from contextlib import nullcontext
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_worker():
|
|
specification = importlib.util.spec_from_file_location(
|
|
"snapform_worker",
|
|
PROJECT_ROOT / "src" / "snapform" / "worker.py",
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class WorkerTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.worker = load_worker()
|
|
|
|
def test_foreground_preserves_source_rgb_and_ben2_alpha(self):
|
|
alpha = SimpleNamespace(size=(2, 1))
|
|
source_rgba = SimpleNamespace(size=(2, 1), applied_alpha=None)
|
|
source_rgba.putalpha = lambda value: setattr(
|
|
source_rgba,
|
|
"applied_alpha",
|
|
value,
|
|
)
|
|
source = SimpleNamespace(convert=lambda _mode: source_rgba)
|
|
foreground_rgba = SimpleNamespace(
|
|
getchannel=lambda _name: alpha,
|
|
)
|
|
foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba)
|
|
|
|
result = self.worker.preserve_source_rgb(source, foreground)
|
|
|
|
self.assertIs(result, source_rgba)
|
|
self.assertIs(result.applied_alpha, alpha)
|
|
|
|
def test_collect_images_is_sorted_and_filters_suffixes(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
(root / "b.PNG").touch()
|
|
(root / "a.jpg").touch()
|
|
(root / "notes.txt").touch()
|
|
|
|
images = self.worker.collect_images(root)
|
|
|
|
self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])
|
|
|
|
def test_ai_tasks_support_cooperative_cancellation(self):
|
|
with self.assertRaises(self.worker.TaskCancelled):
|
|
self.worker.check_cancelled(lambda: True)
|
|
|
|
self.worker.check_cancelled(lambda: False)
|
|
|
|
def test_select_device_supports_mps_and_auto_fallback(self):
|
|
mps_torch = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: False),
|
|
backends=SimpleNamespace(
|
|
mps=SimpleNamespace(is_available=lambda: True),
|
|
),
|
|
)
|
|
mps_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type
|
|
)
|
|
|
|
self.assertEqual(
|
|
self.worker.select_device("mps", mps_torch).type,
|
|
"mps",
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("auto", mps_torch).type,
|
|
"mps",
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("cpu", mps_torch).type,
|
|
"cpu",
|
|
)
|
|
|
|
cuda_torch = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: True),
|
|
)
|
|
cuda_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("auto", cuda_torch).type,
|
|
"cuda",
|
|
)
|
|
|
|
cpu_only_torch = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: False),
|
|
backends=SimpleNamespace(
|
|
mps=SimpleNamespace(is_available=lambda: False),
|
|
),
|
|
)
|
|
cpu_only_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("auto", cpu_only_torch).type,
|
|
"cpu",
|
|
)
|
|
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
|
|
self.worker.select_device("mps", cpu_only_torch)
|
|
|
|
def test_reconstruct_parser_accepts_background_model_directory(self):
|
|
args = self.worker.parser().parse_args(
|
|
[
|
|
"reconstruct",
|
|
"--input",
|
|
"input.png",
|
|
"--model-dir",
|
|
"da3",
|
|
"--output-dir",
|
|
"result",
|
|
"--export-format",
|
|
"mini_npz",
|
|
"--status",
|
|
"status.json",
|
|
"--remove-background-model-dir",
|
|
"ben2",
|
|
"--process-res",
|
|
"1024",
|
|
"--process-res-method",
|
|
"upper_bound_resize",
|
|
"--conf-thresh-percentile",
|
|
"40",
|
|
"--num-max-points",
|
|
"1000",
|
|
"--show-cameras",
|
|
"0",
|
|
"--gs-views-interval",
|
|
"1",
|
|
"--video-frames",
|
|
"16",
|
|
"--device",
|
|
"auto",
|
|
"--use-ray-pose",
|
|
"1",
|
|
"--ref-view-strategy",
|
|
"first",
|
|
]
|
|
)
|
|
self.assertEqual(args.remove_background_model_dir, "ben2")
|
|
|
|
def test_ai_depth_is_converted_for_positive_displacement(self):
|
|
prediction = SimpleNamespace(
|
|
depth=np.array([[[0.25, 1.0], [1.5, 2.0]]]),
|
|
conf=np.array([[[0.1, 0.2], [0.3, 0.4]]]),
|
|
)
|
|
|
|
self.worker.convert_ai_depth_for_displacement(prediction)
|
|
|
|
np.testing.assert_array_equal(
|
|
prediction.depth,
|
|
[[[0.75, 0.0], [-0.5, -1.0]]],
|
|
)
|
|
np.testing.assert_array_equal(
|
|
prediction.conf,
|
|
[[[0.1, 0.2], [0.3, 0.4]]],
|
|
)
|
|
|
|
def test_npz_reference_depth_preserves_prediction_arrays(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
npz_path = Path(directory) / "results.npz"
|
|
depth = np.array([[[1.0, 2.0]]], dtype=np.float32)
|
|
confidence = np.array([[[0.2, 0.8]]], dtype=np.float32)
|
|
np.savez_compressed(npz_path, depth=depth, conf=confidence)
|
|
|
|
self.worker.write_npz_reference_depth(npz_path, 1.5)
|
|
|
|
with np.load(npz_path, allow_pickle=False) as archive:
|
|
np.testing.assert_array_equal(archive["depth"], depth)
|
|
np.testing.assert_array_equal(archive["conf"], confidence)
|
|
self.assertEqual(
|
|
float(archive[self.worker.REFERENCE_DEPTH_KEY]),
|
|
1.5,
|
|
)
|
|
|
|
def test_prepares_a_single_image_file(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
image_path = Path(directory) / "input.png"
|
|
image_path.touch()
|
|
|
|
images, temporary_directory = self.worker.prepare_input(
|
|
image_path,
|
|
16,
|
|
directory,
|
|
)
|
|
|
|
self.assertEqual(images, [image_path])
|
|
self.assertIsNone(temporary_directory)
|
|
|
|
def test_prepares_a_frame_directory(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
(root / "frame-0001.jpg").touch()
|
|
(root / "frame-0002.png").touch()
|
|
(root / "notes.txt").touch()
|
|
|
|
images, temporary_directory = self.worker.prepare_input(
|
|
root,
|
|
16,
|
|
directory,
|
|
)
|
|
|
|
self.assertEqual(
|
|
[path.name for path in images],
|
|
["frame-0001.jpg", "frame-0002.png"],
|
|
)
|
|
self.assertIsNone(temporary_directory)
|
|
|
|
def test_prepares_first_frames_of_a_frame_directory(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
for index in range(1, 6):
|
|
(root / f"frame-{index:04d}.jpg").touch()
|
|
|
|
images, temporary_directory = self.worker.prepare_input(
|
|
root,
|
|
3,
|
|
directory,
|
|
)
|
|
|
|
self.assertEqual(
|
|
[path.name for path in images],
|
|
["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
|
|
)
|
|
self.assertIsNone(temporary_directory)
|
|
|
|
def test_prepares_first_frames_after_offset_in_a_frame_directory(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
for index in range(1, 6):
|
|
(root / f"frame-{index:04d}.jpg").touch()
|
|
|
|
images, temporary_directory = self.worker.prepare_input(
|
|
root,
|
|
3,
|
|
directory,
|
|
start_frame=2,
|
|
)
|
|
|
|
self.assertEqual(
|
|
[path.name for path in images],
|
|
["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"],
|
|
)
|
|
self.assertIsNone(temporary_directory)
|
|
|
|
def test_extract_video_frames_reads_frames_sequentially(self):
|
|
class FakeCapture:
|
|
def __init__(self):
|
|
self.position = 0
|
|
|
|
def isOpened(self):
|
|
return True
|
|
|
|
def read(self):
|
|
if self.position >= 100:
|
|
return False, None
|
|
frame = self.position
|
|
self.position += 1
|
|
return True, frame
|
|
|
|
def release(self):
|
|
return None
|
|
|
|
class FakeCv2:
|
|
def __init__(self):
|
|
self.captures = []
|
|
|
|
def VideoCapture(self, _path):
|
|
capture = FakeCapture()
|
|
self.captures.append(capture)
|
|
return capture
|
|
|
|
def imwrite(self, path, frame):
|
|
Path(path).write_text(str(frame), encoding="utf-8")
|
|
return True
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
fake_cv2 = FakeCv2()
|
|
image_paths = self.worker.extract_video_frames(
|
|
root / "clip.mkv",
|
|
root / "out",
|
|
10,
|
|
fake_cv2,
|
|
)
|
|
|
|
self.assertEqual(
|
|
[path.name for path in image_paths],
|
|
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
|
|
)
|
|
self.assertEqual(len(fake_cv2.captures), 1)
|
|
self.assertEqual(fake_cv2.captures[0].position, 10)
|
|
|
|
def test_extract_video_frames_starts_after_offset(self):
|
|
class FakeCapture:
|
|
def __init__(self):
|
|
self.position = 0
|
|
|
|
def isOpened(self):
|
|
return True
|
|
|
|
def read(self):
|
|
if self.position >= 100:
|
|
return False, None
|
|
frame = self.position
|
|
self.position += 1
|
|
return True, frame
|
|
|
|
def release(self):
|
|
return None
|
|
|
|
class FakeCv2:
|
|
def __init__(self):
|
|
self.captures = []
|
|
|
|
def VideoCapture(self, _path):
|
|
capture = FakeCapture()
|
|
self.captures.append(capture)
|
|
return capture
|
|
|
|
def imwrite(self, path, frame):
|
|
Path(path).write_text(str(frame), encoding="utf-8")
|
|
return True
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
fake_cv2 = FakeCv2()
|
|
image_paths = self.worker.extract_video_frames(
|
|
root / "clip.mkv",
|
|
root / "out",
|
|
10,
|
|
fake_cv2,
|
|
start_frame=71,
|
|
)
|
|
extracted = [
|
|
Path(path).read_text(encoding="utf-8")
|
|
for path in image_paths
|
|
]
|
|
|
|
self.assertEqual(extracted, [str(frame) for frame in range(70, 80)])
|
|
self.assertEqual(fake_cv2.captures[0].position, 80)
|
|
|
|
def test_extract_video_frames_tolerates_short_video_tail(self):
|
|
class ShortCapture:
|
|
def __init__(self):
|
|
self.position = 0
|
|
|
|
def isOpened(self):
|
|
return True
|
|
|
|
def read(self):
|
|
if self.position >= 99:
|
|
return False, None
|
|
frame = self.position
|
|
self.position += 1
|
|
return True, frame
|
|
|
|
def release(self):
|
|
return None
|
|
|
|
class FakeCv2:
|
|
def VideoCapture(self, _path):
|
|
return ShortCapture()
|
|
|
|
def imwrite(self, path, frame):
|
|
Path(path).write_text(str(frame), encoding="utf-8")
|
|
return True
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
image_paths = self.worker.extract_video_frames(
|
|
root / "clip.mkv",
|
|
root / "out",
|
|
10,
|
|
FakeCv2(),
|
|
)
|
|
|
|
self.assertEqual(
|
|
[path.name for path in image_paths],
|
|
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
|
|
)
|
|
|
|
def test_npz_is_complete_requires_readable_archive(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
valid_path = root / "valid.npz"
|
|
np.savez_compressed(
|
|
valid_path,
|
|
depth=np.zeros((4, 4)),
|
|
)
|
|
self.assertTrue(self.worker.npz_is_complete(valid_path))
|
|
|
|
empty_path = root / "empty.npz"
|
|
empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
|
|
self.assertFalse(self.worker.npz_is_complete(empty_path))
|
|
|
|
def test_wait_for_exported_paths_waits_for_complete_npz(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
npz_path = root / "exports" / "mini_npz" / "results.npz"
|
|
npz_path.parent.mkdir(parents=True)
|
|
npz_path.touch()
|
|
with patch.object(
|
|
self.worker,
|
|
"npz_is_complete",
|
|
side_effect=[False, True],
|
|
):
|
|
result = self.worker.wait_for_exported_paths(
|
|
root,
|
|
"mini_npz",
|
|
timeout=2.0,
|
|
)
|
|
self.assertEqual(result, [npz_path])
|
|
|
|
def test_move_exported_file_retries_windows_file_lock(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source = root / "source.npz"
|
|
destination = root / "destination.npz"
|
|
source.touch()
|
|
path_type = type(source)
|
|
original_replace = path_type.replace
|
|
calls = 0
|
|
|
|
def locked_once(path, target):
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
error = PermissionError("file is in use")
|
|
error.winerror = 32
|
|
raise error
|
|
return original_replace(path, target)
|
|
|
|
with patch.object(path_type, "replace", new=locked_once), patch.object(
|
|
self.worker.time,
|
|
"sleep",
|
|
) as sleep:
|
|
result = self.worker.move_exported_file(
|
|
source,
|
|
destination,
|
|
)
|
|
|
|
self.assertEqual(result, destination)
|
|
self.assertTrue(destination.is_file())
|
|
sleep.assert_called_once_with(0.1)
|
|
|
|
def test_rejects_unsupported_input_file(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
input_path = Path(directory) / "input.txt"
|
|
input_path.touch()
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
|
|
self.worker.prepare_input(input_path, 16, directory)
|
|
|
|
def test_download_model_skips_unreachable_endpoint(self):
|
|
calls = []
|
|
fake_hugging_face = ModuleType("huggingface_hub")
|
|
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
|
options
|
|
)
|
|
args = SimpleNamespace(
|
|
repository="depth-anything/DA3-SMALL",
|
|
model_dir="model",
|
|
status="status.json",
|
|
allow_pattern=[],
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{"huggingface_hub": fake_hugging_face},
|
|
), patch.object(
|
|
self.worker,
|
|
"create_download_progress_class",
|
|
return_value=object,
|
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
|
self.worker,
|
|
"hf_endpoint_reachable",
|
|
return_value=False,
|
|
), patch.object(
|
|
self.worker,
|
|
"download_from_modelscope",
|
|
) as mocked_modelscope, patch.object(
|
|
self.worker,
|
|
"_verify_model_complete",
|
|
return_value=None,
|
|
):
|
|
args.model_dir = str(Path(directory) / "model")
|
|
args.status = str(Path(directory) / "status.json")
|
|
self.worker.download_model(args)
|
|
|
|
self.assertEqual(calls, [])
|
|
self.assertEqual(mocked_modelscope.call_count, 1)
|
|
|
|
def test_download_model_refuses_success_when_files_missing(self):
|
|
fake_hugging_face = ModuleType("huggingface_hub")
|
|
fake_hugging_face.snapshot_download = lambda **options: None
|
|
args = SimpleNamespace(
|
|
repository="PramaLLC/BEN2",
|
|
model_dir="model",
|
|
status="status.json",
|
|
allow_pattern=[],
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{"huggingface_hub": fake_hugging_face},
|
|
), patch.object(
|
|
self.worker,
|
|
"create_download_progress_class",
|
|
return_value=object,
|
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
|
self.worker,
|
|
"download_from_modelscope",
|
|
):
|
|
args.model_dir = str(Path(directory) / "model")
|
|
args.status = str(Path(directory) / "status.json")
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"All download sources failed",
|
|
):
|
|
self.worker.download_model(args)
|
|
|
|
def test_download_model_reports_all_failed_sources(self):
|
|
fake_hugging_face = ModuleType("huggingface_hub")
|
|
|
|
def fake_snapshot_download(**options):
|
|
raise OSError("direct down")
|
|
|
|
fake_hugging_face.snapshot_download = fake_snapshot_download
|
|
args = SimpleNamespace(
|
|
repository="depth-anything/DA3-SMALL",
|
|
model_dir="model",
|
|
status="status.json",
|
|
allow_pattern=[],
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{"huggingface_hub": fake_hugging_face},
|
|
), patch.object(
|
|
self.worker,
|
|
"create_download_progress_class",
|
|
return_value=object,
|
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
|
self.worker,
|
|
"download_from_modelscope",
|
|
side_effect=OSError("mirror down"),
|
|
):
|
|
args.model_dir = str(Path(directory) / "model")
|
|
args.status = str(Path(directory) / "status.json")
|
|
with self.assertRaisesRegex(RuntimeError, "modelscope.cn"):
|
|
self.worker.download_model(args)
|
|
|
|
def test_download_from_modelscope_downloads_selected_files(self):
|
|
args = SimpleNamespace(
|
|
repository="depth-anything/DA3-SMALL",
|
|
model_dir="model",
|
|
status="status.json",
|
|
allow_pattern=["config.json"],
|
|
)
|
|
downloaded = []
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.object(
|
|
self.worker,
|
|
"modelscope_file_list",
|
|
return_value=[
|
|
("config.json", 100),
|
|
("model.safetensors", 200),
|
|
(".gitattributes", 10),
|
|
],
|
|
), patch.object(
|
|
self.worker,
|
|
"_download_modelscope_file",
|
|
side_effect=lambda *args, **kwargs: downloaded.append(args[1]),
|
|
), patch.object(self.worker, "write_status"):
|
|
args.model_dir = str(Path(directory) / "model")
|
|
args.status = str(Path(directory) / "status.json")
|
|
self.worker.download_from_modelscope(args, args.model_dir)
|
|
|
|
self.assertEqual(downloaded, ["config.json"])
|
|
|
|
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
|
inference_options = []
|
|
save_options = []
|
|
|
|
class FakeImage:
|
|
size = (1, 1)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def convert(self, mode):
|
|
self.mode = mode
|
|
return self
|
|
|
|
def getchannel(self, _name):
|
|
return self
|
|
|
|
def putalpha(self, _alpha):
|
|
return None
|
|
|
|
def save(self, path, **options):
|
|
save_options.append(options)
|
|
Path(path).touch()
|
|
|
|
class FakeModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_dir):
|
|
self.assertEqual(model_dir, "model")
|
|
return cls()
|
|
|
|
def to(model_self, device):
|
|
self.assertEqual(device.type, "cpu")
|
|
return model_self
|
|
|
|
def eval(self):
|
|
return None
|
|
|
|
def inference(self, image, **options):
|
|
inference_options.append((image, options))
|
|
return FakeImage()
|
|
|
|
fake_torch = ModuleType("torch")
|
|
fake_torch.cuda = SimpleNamespace(
|
|
is_available=lambda: False,
|
|
synchronize=lambda: None,
|
|
)
|
|
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
|
fake_torch.inference_mode = nullcontext
|
|
fake_ben2 = ModuleType("ben2")
|
|
fake_ben2.BEN_Base = FakeModel
|
|
fake_pil = ModuleType("PIL")
|
|
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{
|
|
"torch": fake_torch,
|
|
"ben2": fake_ben2,
|
|
"PIL": fake_pil,
|
|
},
|
|
):
|
|
root = Path(directory)
|
|
input_path = root / "input.png"
|
|
input_path.touch()
|
|
args = SimpleNamespace(
|
|
input=str(input_path),
|
|
model_dir="model",
|
|
output_dir=str(root / "result"),
|
|
status=str(root / "status.json"),
|
|
device="cpu",
|
|
)
|
|
self.worker.remove_background(args)
|
|
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
|
|
|
self.assertTrue((root / "result" / "foreground.png").is_file())
|
|
|
|
self.assertEqual(
|
|
inference_options[0][1],
|
|
{"refine_foreground": False},
|
|
)
|
|
self.assertEqual(
|
|
save_options[0],
|
|
{"format": "PNG", "compress_level": 1},
|
|
)
|
|
self.assertIn("Background removed in", status["message"])
|
|
|
|
def test_ben2_background_removal_processes_each_frame(self):
|
|
inference_options = []
|
|
|
|
class FakeImage:
|
|
size = (1, 1)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def convert(self, mode):
|
|
self.mode = mode
|
|
return self
|
|
|
|
def getchannel(self, _name):
|
|
return self
|
|
|
|
def putalpha(self, _alpha):
|
|
return None
|
|
|
|
def save(self, path, **options):
|
|
Path(path).touch()
|
|
|
|
class FakeModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_dir):
|
|
return cls()
|
|
|
|
def to(model_self, device):
|
|
return model_self
|
|
|
|
def eval(self):
|
|
return None
|
|
|
|
def inference(self, image, **options):
|
|
inference_options.append((image, options))
|
|
return FakeImage()
|
|
|
|
fake_torch = ModuleType("torch")
|
|
fake_torch.cuda = SimpleNamespace(
|
|
is_available=lambda: False,
|
|
synchronize=lambda: None,
|
|
)
|
|
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
|
fake_torch.inference_mode = nullcontext
|
|
fake_ben2 = ModuleType("ben2")
|
|
fake_ben2.BEN_Base = FakeModel
|
|
fake_pil = ModuleType("PIL")
|
|
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{
|
|
"torch": fake_torch,
|
|
"ben2": fake_ben2,
|
|
"PIL": fake_pil,
|
|
},
|
|
):
|
|
root = Path(directory)
|
|
input_directory = root / "frames"
|
|
input_directory.mkdir()
|
|
(input_directory / "frame-0001.jpg").touch()
|
|
(input_directory / "frame-0002.jpg").touch()
|
|
args = SimpleNamespace(
|
|
input=str(input_directory),
|
|
model_dir="model",
|
|
output_dir=str(root / "result"),
|
|
status=str(root / "status.json"),
|
|
device="cpu",
|
|
video_frames=8,
|
|
)
|
|
self.worker.remove_background(args)
|
|
|
|
result_directory = Path(args.output_dir)
|
|
self.assertTrue(
|
|
(result_directory / "foreground.0001.png").is_file()
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "foreground.0002.png").is_file()
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "frame.0001.jpg").is_file()
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "frame.0002.jpg").is_file()
|
|
)
|
|
|
|
self.assertEqual(len(inference_options), 2)
|
|
for _image, options in inference_options:
|
|
self.assertEqual(options, {"refine_foreground": False})
|
|
|
|
def test_remove_background_parser_accepts_video_frames(self):
|
|
args = self.worker.parser().parse_args(
|
|
[
|
|
"remove-background",
|
|
"--input",
|
|
"input.mp4",
|
|
"--model-dir",
|
|
"ben2",
|
|
"--output-dir",
|
|
"result",
|
|
"--status",
|
|
"status.json",
|
|
"--device",
|
|
"auto",
|
|
"--video-frames",
|
|
"24",
|
|
]
|
|
)
|
|
|
|
self.assertEqual(args.command, "remove-background")
|
|
self.assertEqual(args.video_frames, 24)
|
|
|
|
def test_extract_frames_parser_accepts_frame_range(self):
|
|
args = self.worker.parser().parse_args(
|
|
[
|
|
"extract-frames",
|
|
"--input",
|
|
"frames",
|
|
"--output-dir",
|
|
"result",
|
|
"--status",
|
|
"status.json",
|
|
"--video-frames",
|
|
"24",
|
|
"--start-frame",
|
|
"70",
|
|
]
|
|
)
|
|
|
|
self.assertEqual(args.command, "extract-frames")
|
|
self.assertEqual(args.video_frames, 24)
|
|
self.assertEqual(args.start_frame, 70)
|
|
|
|
def test_extract_frames_copies_selected_frames(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_directory = root / "frames"
|
|
input_directory.mkdir()
|
|
for index in range(1, 6):
|
|
(input_directory / f"frame-{index:04d}.jpg").touch()
|
|
args = SimpleNamespace(
|
|
input=str(input_directory),
|
|
output_dir=str(root / "result"),
|
|
status=str(root / "status.json"),
|
|
video_frames=3,
|
|
start_frame=2,
|
|
)
|
|
|
|
self.worker.extract_frames(args)
|
|
|
|
result_directory = Path(args.output_dir)
|
|
self.assertEqual(
|
|
sorted(
|
|
path.name
|
|
for path in result_directory.iterdir()
|
|
if path.is_file()
|
|
),
|
|
["frame.0001.jpg", "frame.0002.jpg", "frame.0003.jpg"],
|
|
)
|
|
status = json.loads(
|
|
Path(args.status).read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(status["progress"], 1.0)
|
|
self.assertIn("3 frame", status["message"])
|
|
|
|
def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
|
|
inference_options = []
|
|
export_calls = []
|
|
mesh_calls = []
|
|
|
|
class FakeModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_dir):
|
|
if model_dir != "model":
|
|
raise AssertionError(f"Unexpected model directory: {model_dir}")
|
|
return cls()
|
|
|
|
def to(self, device):
|
|
if device.type != "cpu":
|
|
raise AssertionError(f"Unexpected device: {device.type}")
|
|
return self
|
|
|
|
def eval(self):
|
|
return None
|
|
|
|
def inference(self, images, **options):
|
|
inference_options.append((images, options))
|
|
return SimpleNamespace(
|
|
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
|
|
)
|
|
|
|
def fake_export(prediction, export_format, export_dir, **options):
|
|
export_calls.append((prediction, export_format, export_dir, options))
|
|
output_directory = Path(export_dir)
|
|
if export_format == "glb":
|
|
(output_directory / "scene.glb").touch()
|
|
elif export_format == "gs_ply":
|
|
ply_directory = output_directory / "gs_ply"
|
|
ply_directory.mkdir()
|
|
(ply_directory / "0000.ply").touch()
|
|
elif export_format == "depth_vis":
|
|
depth_directory = output_directory / "depth_vis"
|
|
depth_directory.mkdir()
|
|
(depth_directory / "0000.jpg").touch()
|
|
else:
|
|
mini_directory = output_directory / "exports" / "mini_npz"
|
|
mini_directory.mkdir(parents=True)
|
|
np.savez_compressed(
|
|
mini_directory / "results.npz",
|
|
depth=np.zeros((14, 14)),
|
|
)
|
|
|
|
fake_torch = ModuleType("torch")
|
|
fake_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type,
|
|
)
|
|
fake_torch.inference_mode = nullcontext
|
|
fake_api = ModuleType("depth_anything_3.api")
|
|
fake_api.DepthAnything3 = FakeModel
|
|
fake_export_module = ModuleType("depth_anything_3.utils.export")
|
|
fake_export_module.export = fake_export
|
|
fake_utils = ModuleType("depth_anything_3.utils")
|
|
fake_utils.__path__ = []
|
|
fake_package = ModuleType("depth_anything_3")
|
|
fake_package.__path__ = []
|
|
|
|
def fake_mesh_export(
|
|
prediction,
|
|
images,
|
|
output_directory,
|
|
confidence_percentile,
|
|
filter_edges,
|
|
max_edge_length,
|
|
):
|
|
mesh_calls.append(
|
|
(
|
|
prediction,
|
|
images,
|
|
confidence_percentile,
|
|
filter_edges,
|
|
max_edge_length,
|
|
)
|
|
)
|
|
output_path = Path(output_directory) / "mesh.glb"
|
|
output_path.touch()
|
|
return output_path
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_path = root / "input.png"
|
|
input_path.touch()
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"torch": fake_torch,
|
|
"depth_anything_3": fake_package,
|
|
"depth_anything_3.api": fake_api,
|
|
"depth_anything_3.utils": fake_utils,
|
|
"depth_anything_3.utils.export": fake_export_module,
|
|
},
|
|
), patch.object(
|
|
self.worker,
|
|
"export_textured_mesh_glb",
|
|
side_effect=fake_mesh_export,
|
|
):
|
|
debug_formats = (
|
|
"glb",
|
|
"gs_ply",
|
|
"depth_vis",
|
|
"mini_npz",
|
|
"mesh",
|
|
)
|
|
for export_format in debug_formats:
|
|
args = SimpleNamespace(
|
|
input=str(input_path),
|
|
video_frames=16,
|
|
output_dir=str(root / export_format),
|
|
export_format=export_format,
|
|
status=str(root / f"{export_format}.json"),
|
|
device="cpu",
|
|
model_dir="model",
|
|
process_res=518,
|
|
process_res_method="lower_bound_resize",
|
|
use_ray_pose=1,
|
|
ref_view_strategy="middle",
|
|
conf_thresh_percentile=35.0,
|
|
num_max_points=12345,
|
|
show_cameras=0,
|
|
gs_views_interval=2,
|
|
mesh_filter_edges=True,
|
|
mesh_max_edge_length=0.2,
|
|
)
|
|
self.worker.reconstruct(args)
|
|
status = json.loads(
|
|
Path(args.status).read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(status["progress"], 1.0)
|
|
self.assertIn(export_format, status["message"])
|
|
result_directory = Path(args.output_dir)
|
|
if export_format == "gs_ply":
|
|
self.assertTrue((result_directory / "0000.ply").is_file())
|
|
self.assertFalse((result_directory / "gs_ply").exists())
|
|
if export_format == "depth_vis":
|
|
self.assertTrue((result_directory / "0000.jpg").is_file())
|
|
self.assertFalse((result_directory / "depth_vis").exists())
|
|
|
|
self.assertEqual(len(inference_options), 5)
|
|
self.assertEqual(len(export_calls), 4)
|
|
self.assertEqual(len(mesh_calls), 1)
|
|
for images, options in inference_options:
|
|
self.assertEqual(images, [str(input_path)])
|
|
self.assertEqual(options["process_res"], 518)
|
|
self.assertEqual(
|
|
options["process_res_method"],
|
|
"lower_bound_resize",
|
|
)
|
|
self.assertTrue(options["use_ray_pose"])
|
|
self.assertEqual(options["ref_view_strategy"], "middle")
|
|
for (
|
|
_prediction,
|
|
export_format,
|
|
_export_dir,
|
|
options,
|
|
) in export_calls:
|
|
expected_depth = (
|
|
-1.0 if export_format == "mini_npz" else 2.0
|
|
)
|
|
np.testing.assert_array_equal(
|
|
_prediction.depth,
|
|
np.full((1, 2, 2), expected_depth, dtype=np.float32),
|
|
)
|
|
if export_format == "glb":
|
|
self.assertEqual(
|
|
options,
|
|
{
|
|
"glb": {
|
|
"conf_thresh_percentile": 35.0,
|
|
"num_max_points": 12345,
|
|
"show_cameras": False,
|
|
"export_depth_vis": False,
|
|
}
|
|
},
|
|
)
|
|
elif export_format == "gs_ply":
|
|
self.assertEqual(
|
|
options,
|
|
{"gs_ply": {"gs_views_interval": 2}},
|
|
)
|
|
else:
|
|
self.assertEqual(options, {})
|
|
self.assertEqual(mesh_calls[0][1], [input_path])
|
|
self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2))
|
|
|
|
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
|
|
self.assertEqual(
|
|
infer_gs_values,
|
|
[False, True, False, False, False],
|
|
)
|
|
|
|
def test_reconstruct_copies_frames_for_multi_frame_input(self):
|
|
class FakeModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_dir):
|
|
return cls()
|
|
|
|
def to(self, device):
|
|
return self
|
|
|
|
def eval(self):
|
|
return None
|
|
|
|
def inference(self, images, **options):
|
|
return SimpleNamespace(
|
|
depth=np.full((2, 2, 2), 2.0, dtype=np.float32),
|
|
)
|
|
|
|
def fake_export(prediction, export_format, export_dir, **options):
|
|
mini_directory = Path(export_dir) / "exports" / "mini_npz"
|
|
mini_directory.mkdir(parents=True)
|
|
np.savez_compressed(
|
|
mini_directory / "results.npz",
|
|
depth=np.zeros((2, 14, 14)),
|
|
)
|
|
|
|
fake_torch = ModuleType("torch")
|
|
fake_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type,
|
|
)
|
|
fake_torch.inference_mode = nullcontext
|
|
fake_api = ModuleType("depth_anything_3.api")
|
|
fake_api.DepthAnything3 = FakeModel
|
|
fake_export_module = ModuleType("depth_anything_3.utils.export")
|
|
fake_export_module.export = fake_export
|
|
fake_utils = ModuleType("depth_anything_3.utils")
|
|
fake_utils.__path__ = []
|
|
fake_package = ModuleType("depth_anything_3")
|
|
fake_package.__path__ = []
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_directory = root / "frames"
|
|
input_directory.mkdir()
|
|
(input_directory / "frame-0001.jpg").touch()
|
|
(input_directory / "frame-0002.jpg").touch()
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"torch": fake_torch,
|
|
"depth_anything_3": fake_package,
|
|
"depth_anything_3.api": fake_api,
|
|
"depth_anything_3.utils": fake_utils,
|
|
"depth_anything_3.utils.export": fake_export_module,
|
|
},
|
|
):
|
|
args = SimpleNamespace(
|
|
input=str(input_directory),
|
|
video_frames=16,
|
|
output_dir=str(root / "result"),
|
|
export_format="mini_npz",
|
|
status=str(root / "status.json"),
|
|
device="cpu",
|
|
model_dir="model",
|
|
process_res=518,
|
|
process_res_method="upper_bound_resize",
|
|
use_ray_pose=1,
|
|
ref_view_strategy="first",
|
|
conf_thresh_percentile=40.0,
|
|
num_max_points=1000,
|
|
show_cameras=0,
|
|
gs_views_interval=1,
|
|
remove_background_model_dir="",
|
|
)
|
|
self.worker.reconstruct(args)
|
|
|
|
result_directory = Path(args.output_dir)
|
|
self.assertTrue(
|
|
(result_directory / "frame.0001.jpg").is_file()
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "frame.0002.jpg").is_file()
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "results.npz").is_file()
|
|
)
|
|
self.assertFalse(
|
|
(result_directory / "exports").exists()
|
|
)
|
|
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
|
self.assertEqual(status["progress"], 1.0)
|
|
|
|
def test_mesh_confidence_uses_percentile(self):
|
|
depth = np.ones((1, 2, 2), dtype=np.float32)
|
|
confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]])
|
|
|
|
threshold = self.worker.confidence_percentile_threshold(
|
|
depth,
|
|
confidence,
|
|
50.0,
|
|
)
|
|
|
|
self.assertEqual(threshold, 2.5)
|
|
|
|
def test_zero_mesh_confidence_disables_confidence_filter(self):
|
|
depth = np.ones((2, 2), dtype=np.float32)
|
|
confidence = np.full((2, 2), np.nan, dtype=np.float32)
|
|
x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0])
|
|
points = np.stack((x, y, depth), axis=-1)
|
|
|
|
threshold = self.worker.confidence_percentile_threshold(
|
|
depth,
|
|
confidence,
|
|
0.0,
|
|
)
|
|
faces = self.worker.mesh_triangle_faces(
|
|
points,
|
|
depth,
|
|
confidence,
|
|
threshold,
|
|
filter_edges=False,
|
|
max_edge_length=0.0,
|
|
)
|
|
|
|
self.assertIsNone(threshold)
|
|
self.assertEqual(len(faces), 2)
|
|
|
|
def test_mesh_filters_triangles_independently(self):
|
|
depth = np.ones((2, 2), dtype=np.float32)
|
|
confidence = np.ones((2, 2), dtype=np.float32)
|
|
confidence[1, 0] = 0.0
|
|
x, y = np.meshgrid(
|
|
np.arange(2, dtype=np.float32),
|
|
np.arange(2, dtype=np.float32),
|
|
)
|
|
points = np.stack((x, y, depth), axis=-1)
|
|
|
|
faces = self.worker.mesh_triangle_faces(
|
|
points,
|
|
depth,
|
|
confidence,
|
|
confidence_threshold=0.5,
|
|
filter_edges=False,
|
|
max_edge_length=0.0,
|
|
)
|
|
|
|
np.testing.assert_array_equal(faces, [[0, 3, 1]])
|
|
|
|
def test_mesh_background_mask_combines_with_confidence(self):
|
|
depth = np.ones((2, 3), dtype=np.float32)
|
|
confidence = np.ones((2, 3), dtype=np.float32)
|
|
x, y = np.meshgrid(
|
|
np.arange(3, dtype=np.float32),
|
|
np.arange(2, dtype=np.float32),
|
|
)
|
|
points = np.stack((x, y, depth), axis=-1)
|
|
foreground = np.ones((2, 3), dtype=bool)
|
|
foreground[:, 0] = False
|
|
|
|
faces = self.worker.mesh_triangle_faces(
|
|
points,
|
|
depth,
|
|
confidence,
|
|
confidence_threshold=None,
|
|
filter_edges=False,
|
|
max_edge_length=0.0,
|
|
vertex_mask=foreground,
|
|
)
|
|
|
|
np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])
|
|
|
|
def test_depth_edge_filter_ignores_smooth_slope(self):
|
|
smooth = np.tile(
|
|
np.linspace(2.0, 2.04, 5, dtype=np.float32),
|
|
(3, 1),
|
|
)
|
|
discontinuous = smooth.copy()
|
|
discontinuous[:, 3:] += 1.0
|
|
|
|
self.assertFalse(self.worker.depth_edge_mask(smooth).any())
|
|
self.assertTrue(self.worker.depth_edge_mask(discontinuous)[:, 2:4].all())
|
|
|
|
def test_higher_edge_culling_allows_less_stretch(self):
|
|
limits = [
|
|
self.worker.edge_culling_stretch_limit(value)
|
|
for value in (0.0, 50.0, 100.0)
|
|
]
|
|
|
|
self.assertEqual(limits, [10000.0, 2.0, 1.0])
|
|
|
|
def test_stretched_edge_culling_uses_camera_space_edge_ratio(self):
|
|
flat_points = np.asarray(
|
|
[
|
|
[[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]],
|
|
[[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
stretched_points = flat_points.copy()
|
|
stretched_points[0, 1] = (2.0, 0.0, 2.0)
|
|
triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32)
|
|
|
|
flat_mask = self.worker.triangle_edge_stretch_mask(
|
|
flat_points,
|
|
triangles,
|
|
1.0,
|
|
)
|
|
loose_mask = self.worker.triangle_edge_stretch_mask(
|
|
stretched_points,
|
|
triangles,
|
|
2.0,
|
|
)
|
|
strict_mask = self.worker.triangle_edge_stretch_mask(
|
|
stretched_points,
|
|
triangles,
|
|
1.0,
|
|
)
|
|
|
|
self.assertTrue(flat_mask.all())
|
|
self.assertTrue(loose_mask.all())
|
|
self.assertFalse(strict_mask.all())
|
|
|
|
def test_mesh_faces_apply_confidence_threshold_and_face_direction(self):
|
|
points = np.array(
|
|
[
|
|
[[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [2.0, 0.0, 1.0]],
|
|
[[0.0, 1.0, 1.0], [1.0, 1.0, 1.0], [2.0, 1.0, 1.0]],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
depth = np.ones((2, 3), dtype=np.float32)
|
|
confidence = np.ones((2, 3), dtype=np.float32)
|
|
confidence[:, 0] = 0.0
|
|
|
|
faces = self.worker.mesh_triangle_faces(
|
|
points,
|
|
depth,
|
|
confidence,
|
|
confidence_threshold=0.5,
|
|
filter_edges=False,
|
|
max_edge_length=0.0,
|
|
)
|
|
|
|
np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]])
|
|
vertices = points.reshape(-1, 3)[faces[0]]
|
|
normal = np.cross(
|
|
vertices[1] - vertices[0],
|
|
vertices[2] - vertices[0],
|
|
)
|
|
self.assertLess(normal[2], 0.0)
|
|
|
|
def test_mesh_percentile_reduces_face_count(self):
|
|
depth = np.ones((3, 3), dtype=np.float32)
|
|
confidence = np.arange(1, 10, dtype=np.float32).reshape(3, 3)
|
|
x, y = np.meshgrid(
|
|
np.arange(3, dtype=np.float32),
|
|
np.arange(3, dtype=np.float32),
|
|
)
|
|
points = np.stack((x, y, depth), axis=-1)
|
|
face_counts = []
|
|
for percentile in (0.0, 50.0, 100.0):
|
|
threshold = self.worker.confidence_percentile_threshold(
|
|
depth,
|
|
confidence,
|
|
percentile,
|
|
)
|
|
faces = self.worker.mesh_triangle_faces(
|
|
points,
|
|
depth,
|
|
confidence,
|
|
confidence_threshold=threshold,
|
|
filter_edges=False,
|
|
max_edge_length=0.0,
|
|
)
|
|
face_counts.append(len(faces))
|
|
|
|
self.assertEqual(face_counts, [8, 2, 0])
|
|
|
|
def test_mesh_compaction_removes_unreferenced_vertices(self):
|
|
vertices = np.arange(18, dtype=np.float32).reshape(6, 3)
|
|
faces = np.asarray([[1, 3, 5]], dtype=np.int32)
|
|
|
|
compact_vertices, compact_faces = self.worker.compact_mesh(
|
|
vertices,
|
|
faces,
|
|
)
|
|
|
|
np.testing.assert_array_equal(
|
|
compact_vertices,
|
|
vertices[[1, 3, 5]],
|
|
)
|
|
np.testing.assert_array_equal(compact_faces, [[0, 1, 2]])
|
|
|
|
def test_image_mesh_data_keeps_camera_geometry_and_uv(self):
|
|
class FakeImage:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def convert(self, mode):
|
|
self.mode = mode
|
|
return self
|
|
|
|
def save(self, path, **_options):
|
|
Path(path).write_bytes(b"texture")
|
|
|
|
fake_pil = ModuleType("PIL")
|
|
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
|
depth = np.full((1, 3, 3), 2.0, dtype=np.float32)
|
|
confidence = np.ones((1, 3, 3), dtype=np.float32)
|
|
intrinsics = np.asarray(
|
|
[[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]]],
|
|
dtype=np.float32,
|
|
)
|
|
prediction = SimpleNamespace(
|
|
depth=depth,
|
|
conf=confidence,
|
|
intrinsics=intrinsics,
|
|
)
|
|
progress = []
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{"PIL": fake_pil},
|
|
), patch.object(
|
|
self.worker,
|
|
"compact_mesh",
|
|
side_effect=lambda vertices, faces: (
|
|
vertices,
|
|
faces,
|
|
),
|
|
):
|
|
result_path, texture_path = self.worker.export_image_mesh_data(
|
|
prediction,
|
|
Path(directory) / "input.png",
|
|
Path(directory) / "result",
|
|
progress_callback=lambda step, total, label: progress.append(
|
|
(step, total, label)
|
|
),
|
|
)
|
|
with np.load(result_path) as result:
|
|
self.assertEqual(result["vertices"].shape, (9, 3))
|
|
self.assertEqual(result["faces"].shape, (8, 3))
|
|
self.assertEqual(result["uv"].shape, (9, 2))
|
|
self.assertEqual(float(result["reference_depth"]), 2.0)
|
|
self.assertNotIn("postprocess_method", result.files)
|
|
self.assertEqual(result["point_depth"].shape, (9,))
|
|
self.assertEqual(
|
|
result["point_confidence_percentile"].shape,
|
|
(9,),
|
|
)
|
|
self.assertEqual(
|
|
result["face_edge_stretch"].shape,
|
|
(8,),
|
|
)
|
|
self.assertEqual(result["face_background"].shape, (8,))
|
|
np.testing.assert_allclose(result["uv"][0], [0.0, 1.0])
|
|
np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0])
|
|
self.assertTrue(texture_path.is_file())
|
|
self.assertEqual([item[0] for item in progress], list(range(1, 7)))
|
|
self.assertTrue(all(item[1] == 6 for item in progress))
|
|
|
|
def test_prediction_attributes_keep_debug_structure(self):
|
|
prediction = SimpleNamespace(
|
|
depth=np.ones((1, 2, 3), dtype=np.float32),
|
|
conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3),
|
|
intrinsics=np.eye(3, dtype=np.float32)[None],
|
|
extrinsics=np.eye(4, dtype=np.float32)[None],
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
output_path = self.worker.export_prediction_attributes(
|
|
prediction,
|
|
directory,
|
|
)
|
|
with np.load(output_path) as attributes:
|
|
self.assertEqual(attributes["depth"].shape, (1, 2, 3))
|
|
self.assertEqual(attributes["conf"].shape, (1, 2, 3))
|
|
self.assertEqual(
|
|
attributes["confidence_percentile"].shape,
|
|
(1, 2, 3),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|