Implement Omoo DA3 functionality with property group, runtime, worker, and status bar
- Added properties for Omoo DA3 settings in properties.py, including input path, video frames, model selection, and export format. - Created runtime management in runtime.py, handling model downloads, job paths, and command execution for reconstruction. - Developed status bar functionality in status_bar.py to display job progress and status. - Implemented worker.py for image and video processing, including frame extraction and model inference. - Added comprehensive unit tests covering addon registration, model handling, export processes, and UI interactions. - Included a lock file for dependency management with uv.
This commit is contained in:
194
tests/test_worker.py
Normal file
194
tests/test_worker.py
Normal file
@@ -0,0 +1,194 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_worker():
|
||||
specification = importlib.util.spec_from_file_location(
|
||||
"omoo_da3_worker",
|
||||
PROJECT_ROOT / "src" / "omoo_da3" / "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_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_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_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_uses_official_da3_export_parameters_for_all_four_formats(self):
|
||||
inference_options = []
|
||||
export_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()
|
||||
|
||||
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)
|
||||
(mini_directory / "results.npz").touch()
|
||||
|
||||
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_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,
|
||||
},
|
||||
):
|
||||
for export_format in self.worker.EXPORT_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,
|
||||
)
|
||||
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"])
|
||||
|
||||
self.assertEqual(len(inference_options), 4)
|
||||
self.assertEqual(len(export_calls), 4)
|
||||
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:
|
||||
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, {})
|
||||
|
||||
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
|
||||
self.assertEqual(infer_gs_values, [False, True, False, False])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user