icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add ONNX model inference implementations for MoGe-2, BEN2, and DA3

Browse Source 
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing.
- Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing.
- Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Developed unit tests for server functionalities in `test_server.py`, including model caching and job management.
- Enhanced job server to handle frame extraction without requiring models.
This commit is contained in:
马南
2026-08-17 10:37:55 +08:00
parent b3a07cde32
commit c5b77c5c85
58 changed files with 5840 additions and 9931 deletions
Download Patch File Download Diff File Expand all files Collapse all files

502
tests/test_worker.py
View File

@@ -1,7 +1,5 @@
import importlib.util
import json
import os
import shutil
import sys
import tempfile
import unittest
@@ -16,14 +14,13 @@ 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" / "ai_service" / "jobs.py",
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
def load_server_module(name):
addon_directory = PROJECT_ROOT / "src" / "snapform"
sys.path.insert(0, str(addon_directory))
try:
return __import__(f"server.{name}", fromlist=[name])
finally:
sys.path.remove(str(addon_directory))
def _make_fake_av(frame_count):
@@ -69,10 +66,54 @@ def _make_fake_av(frame_count):
return FakeAv()
class FakeModels:
def __init__(self, directory="model"):
self.model_directory = Path(directory)
def directory(self, _model):
return self.model_directory
def get_ben2(self, _model_directory, _device):
from server.onnx_ben2 import create_session
return create_session(self.model_directory, "cpu"), 0.0
def get_da3(self, _model_directory, _device):
return object()
class FakeJobContext:
def __init__(self, directory, job_type, models=None):
self.directory = Path(directory)
self.job_type = job_type
self.models = models or FakeModels()
self.status = {}
def resource(self, name):
if name != "model_manager":
raise KeyError(name)
return self.models
def progress(self, progress, message, **details):
self.status = {
"progress": progress,
"message": message,
**details,
}
def check_cancelled(self):
return None
class WorkerTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.worker = load_worker()
cls.media = load_server_module("media")
cls.geometry = load_server_module("geometry")
cls.background = load_server_module("background")
cls.moge2 = load_server_module("moge2")
cls.da3 = load_server_module("da3")
cls.download = load_server_module("model_download")
def test_foreground_preserves_source_rgb_and_ben2_alpha(self):
alpha = SimpleNamespace(size=(2, 1))
@@ -88,7 +129,7 @@ class WorkerTest(unittest.TestCase):
)
foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba)
result = self.worker.preserve_source_rgb(source, foreground)
result = self.background.preserve_source_rgb(source, foreground)
self.assertIs(result, source_rgba)
self.assertIs(result.applied_alpha, alpha)
@@ -100,62 +141,22 @@ class WorkerTest(unittest.TestCase):
(root / "a.jpg").touch()
(root / "notes.txt").touch()
images = self.worker.collect_images(root)
images = self.media.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)
with self.assertRaises(self.media.TaskCancelled):
self.media.check_cancelled(lambda: True)
self.worker.check_cancelled(lambda: False)
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")
self.media.check_cancelled(lambda: False)
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(
images, temporary_directory = self.media.prepare_input(
image_path,
16,
directory,
@@ -171,7 +172,7 @@ class WorkerTest(unittest.TestCase):
(root / "frame-0002.png").touch()
(root / "notes.txt").touch()
images, temporary_directory = self.worker.prepare_input(
images, temporary_directory = self.media.prepare_input(
root,
16,
directory,
@@ -189,7 +190,7 @@ class WorkerTest(unittest.TestCase):
for index in range(1, 6):
(root / f"frame-{index:04d}.jpg").touch()
images, temporary_directory = self.worker.prepare_input(
images, temporary_directory = self.media.prepare_input(
root,
3,
directory,
@@ -207,7 +208,7 @@ class WorkerTest(unittest.TestCase):
for index in range(1, 6):
(root / f"frame-{index:04d}.jpg").touch()
images, temporary_directory = self.worker.prepare_input(
images, temporary_directory = self.media.prepare_input(
root,
3,
directory,
@@ -224,7 +225,7 @@ class WorkerTest(unittest.TestCase):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
fake_av = _make_fake_av(100)
image_paths = self.worker.extract_video_frames(
image_paths = self.media.extract_video_frames(
root / "clip.mkv",
root / "out",
10,
@@ -242,7 +243,7 @@ class WorkerTest(unittest.TestCase):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
fake_av = _make_fake_av(100)
image_paths = self.worker.extract_video_frames(
image_paths = self.media.extract_video_frames(
root / "clip.mkv",
root / "out",
10,
@@ -260,7 +261,7 @@ class WorkerTest(unittest.TestCase):
def test_extract_video_frames_tolerates_short_video_tail(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
image_paths = self.worker.extract_video_frames(
image_paths = self.media.extract_video_frames(
root / "clip.mkv",
root / "out",
10,
@@ -280,11 +281,11 @@ class WorkerTest(unittest.TestCase):
valid_path,
depth=np.zeros((4, 4)),
)
self.assertTrue(self.worker.npz_is_complete(valid_path))
self.assertTrue(self.media.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))
self.assertFalse(self.media.npz_is_complete(empty_path))
def test_wait_for_exported_paths_waits_for_complete_npz(self):
with tempfile.TemporaryDirectory() as directory:
@@ -293,11 +294,11 @@ class WorkerTest(unittest.TestCase):
npz_path.parent.mkdir(parents=True)
npz_path.touch()
with patch.object(
self.worker,
self.media,
"npz_is_complete",
side_effect=[False, True],
):
result = self.worker.wait_for_exported_paths(
result = self.media.wait_for_exported_paths(
root,
"mini_npz",
timeout=2.0,
@@ -324,10 +325,10 @@ class WorkerTest(unittest.TestCase):
return original_replace(path, target)
with patch.object(path_type, "replace", new=locked_once), patch.object(
self.worker.time,
self.media.time,
"sleep",
) as sleep:
result = self.worker.move_exported_file(
result = self.media.move_exported_file(
source,
destination,
)
@@ -342,7 +343,7 @@ class WorkerTest(unittest.TestCase):
input_path.touch()
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
self.worker.prepare_input(input_path, 16, directory)
self.media.prepare_input(input_path, 16, directory)
def test_download_model_uses_actual_request_without_preflight(self):
calls = []
@@ -353,7 +354,7 @@ class WorkerTest(unittest.TestCase):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
progress=lambda *_args, **_kwargs: None,
allow_pattern=[],
)
@@ -361,17 +362,16 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
self.download,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
self.download,
"_verify_model_complete",
return_value=None,
):
args.model_dir = str(Path(directory) / "models" / "DA3")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.download.download_model(args)
self.assertEqual(len(calls), 1)
@@ -381,7 +381,7 @@ class WorkerTest(unittest.TestCase):
args = SimpleNamespace(
repository="PramaLLC/BEN2",
model_dir="model",
status="status.json",
progress=lambda *_args, **_kwargs: None,
allow_pattern=[],
)
@@ -389,47 +389,45 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
self.download,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
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)
self.download.download_model(args)
def test_download_model_reports_all_failed_sources(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
progress=lambda *_args, **_kwargs: None,
allow_pattern=[],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
)
with tempfile.TemporaryDirectory() as directory, patch.object(
self.worker,
self.download,
"download_from_hf",
side_effect=OSError("direct down"),
), patch.object(
self.worker,
self.download,
"download_from_r2",
side_effect=OSError("mirror down"),
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
self.worker.download_model(args)
self.download.download_model(args)
def test_download_model_prefers_direct_hf_then_r2(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
progress=lambda *_args, **_kwargs: None,
allow_pattern=["config.json"],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
@@ -444,83 +442,28 @@ class WorkerTest(unittest.TestCase):
order.append("r2")
with tempfile.TemporaryDirectory() as directory, patch.object(
self.worker,
self.download,
"download_from_hf",
side_effect=fake_download_from_hf,
), patch.object(
self.worker,
self.download,
"download_from_r2",
side_effect=fake_download_from_r2,
), patch.object(
self.worker,
self.download,
"_verify_model_complete",
return_value=None,
):
args.model_dir = str(Path(directory) / "models" / "DA3")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.download.download_model(args)
self.assertEqual(order, ["hf", "r2"])
def test_download_declared_skips_hf_when_r2_preferred(self):
root = PROJECT_ROOT / ".tmp-models-from-r2"
model_dir = root / "models" / "DA3"
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir=str(model_dir),
status="status.json",
allow_pattern=[],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
)
order = []
try:
self.worker._mark_models_from_r2(model_dir)
with patch.object(
self.worker,
"download_from_hf",
side_effect=lambda *a, **k: order.append("hf"),
), patch.object(
self.worker,
"download_from_r2",
side_effect=lambda *a, **k: order.append("r2"),
), patch.object(
self.worker,
"_verify_model_complete",
return_value=None,
), patch.object(self.worker, "write_status"):
self.worker._download_declared_or_fail(
args, model_dir, "DA3 Small", 0.0, 1.0, {}
)
finally:
shutil.rmtree(root, ignore_errors=True)
self.assertEqual(order, ["r2"])
def test_models_from_r2_roundtrip_preserves_fields(self):
root = PROJECT_ROOT / ".tmp-models-from-r2"
model_dir = root / "models" / "DA3"
runtime_state = root / "runtime.json"
try:
runtime_state.parent.mkdir(parents=True, exist_ok=True)
runtime_state.write_text(
json.dumps({"keep": "value"}),
encoding="utf-8",
)
self.assertFalse(self.worker._models_from_r2(model_dir))
self.worker._mark_models_from_r2(model_dir)
self.assertTrue(self.worker._models_from_r2(model_dir))
payload = json.loads(runtime_state.read_text(encoding="utf-8"))
self.assertEqual(payload["keep"], "value")
self.assertTrue(payload["models_from_r2"])
finally:
shutil.rmtree(root, ignore_errors=True)
def test_download_from_r2_downloads_declared_files(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
status="status.json",
progress=lambda *_args, **_kwargs: None,
allow_pattern=["config.json"],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
@@ -528,13 +471,12 @@ class WorkerTest(unittest.TestCase):
downloaded = []
with tempfile.TemporaryDirectory() as directory, patch.object(
self.worker,
self.download,
"_download_file_from_url",
side_effect=lambda *args, **kwargs: downloaded.append(args[1].name),
), patch.object(self.worker, "write_status"):
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_from_r2(args, args.model_dir)
self.download.download_from_r2(args, args.model_dir)
self.assertEqual(downloaded, ["config.json"])
@@ -589,7 +531,7 @@ class WorkerTest(unittest.TestCase):
)
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("ben2_onnx")
fake_ben2 = ModuleType("server.onnx_ben2")
fake_ben2.create_session = lambda _model_dir, _device: object()
fake_ben2.infer = lambda _model, image: (
inference_options.append((image, {})) or FakeImage()
@@ -601,22 +543,18 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{
"torch": fake_torch,
"ben2_onnx": fake_ben2,
"server.onnx_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",
context = FakeJobContext(root / "result", "remove-background")
self.background.remove_background(
context,
{"input": str(input_path), "model": "BEN2", "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())
@@ -628,7 +566,7 @@ class WorkerTest(unittest.TestCase):
save_options[0],
{"format": "PNG", "compress_level": 1},
)
self.assertIn("Background removed in", status["message"])
self.assertIn("Background removed in", context.status["message"])
def test_ben2_background_removal_processes_each_frame(self):
inference_options = []
@@ -677,7 +615,7 @@ class WorkerTest(unittest.TestCase):
)
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("ben2_onnx")
fake_ben2 = ModuleType("server.onnx_ben2")
fake_ben2.create_session = lambda _model_dir, _device: object()
fake_ben2.infer = lambda _model, image: (
inference_options.append((image, {})) or FakeImage()
@@ -689,7 +627,7 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{
"torch": fake_torch,
"ben2_onnx": fake_ben2,
"server.onnx_ben2": fake_ben2,
"PIL": fake_pil,
},
):
@@ -698,17 +636,18 @@ class WorkerTest(unittest.TestCase):
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,
context = FakeJobContext(root / "result", "remove-background")
self.background.remove_background(
context,
{
"input": str(input_directory),
"model": "BEN2",
"device": "cpu",
"video_frames": 8,
},
)
self.worker.remove_background(args)
result_directory = Path(args.output_dir)
result_directory = context.directory
self.assertTrue(
(result_directory / "foreground.0001.png").is_file()
)
@@ -726,49 +665,6 @@ class WorkerTest(unittest.TestCase):
for _image, options in inference_options:
self.assertEqual(options, {})
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)
@@ -776,17 +672,17 @@ class WorkerTest(unittest.TestCase):
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,
context = FakeJobContext(root / "result", "extract-frames")
self.media.extract_frames(
context,
{
"input": str(input_directory),
"video_frames": 3,
"start_frame": 2,
},
)
self.worker.extract_frames(args)
result_directory = Path(args.output_dir)
result_directory = context.directory
self.assertEqual(
sorted(
path.name
@@ -795,11 +691,8 @@ class WorkerTest(unittest.TestCase):
),
["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"])
self.assertEqual(context.status["progress"], 1.0)
self.assertIn("3 frame", context.status["message"])
def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
inference_options = []
@@ -812,7 +705,7 @@ class WorkerTest(unittest.TestCase):
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
)
fake_da3 = ModuleType("da3_onnx")
fake_da3 = ModuleType("server.onnx_da3")
fake_da3.create_session = lambda _model_dir, _device: object()
def fake_infer(_session, images, **options):
@@ -820,7 +713,7 @@ class WorkerTest(unittest.TestCase):
return prediction
fake_da3.infer = fake_infer
fake_runtime = ModuleType("onnx_runtime")
fake_runtime = ModuleType("server.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
def fake_mesh_export(
@@ -865,23 +758,23 @@ class WorkerTest(unittest.TestCase):
with patch.dict(
sys.modules,
{
"da3_onnx": fake_da3,
"onnx_runtime": fake_runtime,
"server.onnx_da3": fake_da3,
"server.onnx_runtime": fake_runtime,
},
), patch.object(
self.worker,
self.da3,
"export_textured_mesh_glb",
side_effect=fake_mesh_export,
), patch.object(
self.worker,
self.da3,
"export_da3_point_cloud",
side_effect=fake_point_cloud,
), patch.object(
self.worker,
self.da3,
"export_da3_depth_preview",
side_effect=fake_depth_preview,
), patch.object(
self.worker,
self.da3,
"export_da3_mini_npz",
side_effect=fake_mini_npz,
):
@@ -892,32 +785,31 @@ class WorkerTest(unittest.TestCase):
"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,
context = FakeJobContext(
root / export_format,
"debug-da3",
)
self.worker.reconstruct(args)
status = json.loads(
Path(args.status).read_text(encoding="utf-8")
self.da3.generate_da3(
context,
{
"input": str(input_path),
"model": "DA3",
"video_frames": 16,
"export_format": export_format,
"device": "cpu",
"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,
"mesh_filter_edges": True,
"mesh_max_edge_length": 0.2,
},
)
self.assertEqual(status["progress"], 1.0)
self.assertIn(export_format, status["message"])
result_directory = Path(args.output_dir)
self.assertEqual(context.status["progress"], 1.0)
self.assertIn(export_format, context.status["message"])
result_directory = context.directory
if export_format == "depth_vis":
self.assertTrue((result_directory / "0000.jpg").is_file())
self.assertFalse((result_directory / "depth_vis").exists())
@@ -938,7 +830,7 @@ class WorkerTest(unittest.TestCase):
def test_reconstruct_copies_frames_for_multi_frame_input(self):
inference_inputs = []
fake_da3 = ModuleType("da3_onnx")
fake_da3 = ModuleType("server.onnx_da3")
fake_da3.create_session = lambda _model_dir, _device: object()
def fake_infer(_session, images, **_options):
@@ -952,7 +844,7 @@ class WorkerTest(unittest.TestCase):
)
fake_da3.infer = fake_infer
fake_runtime = ModuleType("onnx_runtime")
fake_runtime = ModuleType("server.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
with tempfile.TemporaryDirectory() as directory:
@@ -964,32 +856,33 @@ class WorkerTest(unittest.TestCase):
with patch.dict(
sys.modules,
{
"da3_onnx": fake_da3,
"onnx_runtime": fake_runtime,
"server.onnx_da3": fake_da3,
"server.onnx_runtime": fake_runtime,
},
):
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="",
inference_mode="per_frame",
context = FakeJobContext(
root / "result",
"generate-da3-depth",
)
self.da3.generate_da3(
context,
{
"input": str(input_directory),
"model": "DA3",
"video_frames": 16,
"export_format": "mini_npz",
"device": "cpu",
"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,
"inference_mode": "per_frame",
},
)
self.worker.reconstruct(args)
result_directory = Path(args.output_dir)
result_directory = context.directory
self.assertTrue(
(result_directory / "frame.0001.jpg").is_file()
)
@@ -1006,14 +899,13 @@ class WorkerTest(unittest.TestCase):
self.assertFalse(
(result_directory / "exports").exists()
)
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
self.assertEqual(status["progress"], 1.0)
self.assertEqual(context.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(
threshold = self.geometry.confidence_percentile_threshold(
depth,
confidence,
50.0,
@@ -1027,12 +919,12 @@ class WorkerTest(unittest.TestCase):
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(
threshold = self.geometry.confidence_percentile_threshold(
depth,
confidence,
0.0,
)
faces = self.worker.mesh_triangle_faces(
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
@@ -1054,7 +946,7 @@ class WorkerTest(unittest.TestCase):
)
points = np.stack((x, y, depth), axis=-1)
faces = self.worker.mesh_triangle_faces(
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
@@ -1076,7 +968,7 @@ class WorkerTest(unittest.TestCase):
foreground = np.ones((2, 3), dtype=bool)
foreground[:, 0] = False
faces = self.worker.mesh_triangle_faces(
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
@@ -1096,12 +988,12 @@ class WorkerTest(unittest.TestCase):
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())
self.assertFalse(self.geometry.depth_edge_mask(smooth).any())
self.assertTrue(self.geometry.depth_edge_mask(discontinuous)[:, 2:4].all())
def test_higher_edge_culling_allows_less_stretch(self):
limits = [
self.worker.edge_culling_stretch_limit(value)
self.geometry.edge_culling_stretch_limit(value)
for value in (0.0, 50.0, 100.0)
]
@@ -1119,17 +1011,17 @@ class WorkerTest(unittest.TestCase):
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_mask = self.geometry.triangle_edge_stretch_mask(
flat_points,
triangles,
1.0,
)
loose_mask = self.worker.triangle_edge_stretch_mask(
loose_mask = self.geometry.triangle_edge_stretch_mask(
stretched_points,
triangles,
2.0,
)
strict_mask = self.worker.triangle_edge_stretch_mask(
strict_mask = self.geometry.triangle_edge_stretch_mask(
stretched_points,
triangles,
1.0,
@@ -1151,7 +1043,7 @@ class WorkerTest(unittest.TestCase):
confidence = np.ones((2, 3), dtype=np.float32)
confidence[:, 0] = 0.0
faces = self.worker.mesh_triangle_faces(
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
@@ -1178,12 +1070,12 @@ class WorkerTest(unittest.TestCase):
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(
threshold = self.geometry.confidence_percentile_threshold(
depth,
confidence,
percentile,
)
faces = self.worker.mesh_triangle_faces(
faces = self.geometry.mesh_triangle_faces(
points,
depth,
confidence,
@@ -1199,7 +1091,7 @@ class WorkerTest(unittest.TestCase):
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(
compact_vertices, compact_faces = self.geometry.compact_mesh(
vertices,
faces,
)
@@ -1220,7 +1112,7 @@ class WorkerTest(unittest.TestCase):
dtype=bool,
)
filtered = self.worker.remove_background_faces(faces, foreground)
filtered = self.geometry.remove_background_faces(faces, foreground)
np.testing.assert_array_equal(filtered, faces[2:])
@@ -1258,14 +1150,14 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{"PIL": fake_pil},
), patch.object(
self.worker,
self.geometry,
"compact_mesh",
side_effect=lambda vertices, faces: (
vertices,
faces,
),
):
result_path, texture_path = self.worker.export_image_mesh_data(
result_path, texture_path = self.geometry.export_image_mesh_data(
prediction,
Path(directory) / "input.png",
Path(directory) / "result",
@@ -1313,7 +1205,7 @@ class WorkerTest(unittest.TestCase):
)
with tempfile.TemporaryDirectory() as directory:
output_path = self.worker.export_prediction_attributes(
output_path = self.geometry.export_prediction_attributes(
prediction,
directory,
)