From 9586b354403b6b1d377d416040fa9c653d7030e4 Mon Sep 17 00:00:00 2001 From: nan Date: Sun, 2 Aug 2026 00:04:31 +0800 Subject: [PATCH] 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. --- .gitignore | 8 + AGENTS.md | 25 +++ LICENSE | 8 + README.md | 154 ++++++++++++++ pack.py | 32 +++ pyproject.toml | 11 + src/omoo_da3/__init__.py | 41 ++++ src/omoo_da3/blender_manifest.toml | 23 +++ src/omoo_da3/bootstrap_runtime.py | 203 ++++++++++++++++++ src/omoo_da3/models.py | 81 ++++++++ src/omoo_da3/operators.py | 283 +++++++++++++++++++++++++ src/omoo_da3/panel.py | 112 ++++++++++ src/omoo_da3/preferences.py | 52 +++++ src/omoo_da3/properties.py | 146 +++++++++++++ src/omoo_da3/runtime.py | 198 ++++++++++++++++++ src/omoo_da3/status_bar.py | 51 +++++ src/omoo_da3/worker.py | 304 +++++++++++++++++++++++++++ tests/test_addon.py | 322 +++++++++++++++++++++++++++++ tests/test_bootstrap_runtime.py | 63 ++++++ tests/test_manifest.py | 25 +++ tests/test_models.py | 45 ++++ tests/test_pack.py | 39 ++++ tests/test_worker.py | 194 +++++++++++++++++ uv.lock | 8 + 24 files changed, 2428 insertions(+) create mode 100644 .gitignore create mode 100644 AGENTS.md create mode 100644 LICENSE create mode 100644 README.md create mode 100644 pack.py create mode 100644 pyproject.toml create mode 100644 src/omoo_da3/__init__.py create mode 100644 src/omoo_da3/blender_manifest.toml create mode 100644 src/omoo_da3/bootstrap_runtime.py create mode 100644 src/omoo_da3/models.py create mode 100644 src/omoo_da3/operators.py create mode 100644 src/omoo_da3/panel.py create mode 100644 src/omoo_da3/preferences.py create mode 100644 src/omoo_da3/properties.py create mode 100644 src/omoo_da3/runtime.py create mode 100644 src/omoo_da3/status_bar.py create mode 100644 src/omoo_da3/worker.py create mode 100644 tests/test_addon.py create mode 100644 tests/test_bootstrap_runtime.py create mode 100644 tests/test_manifest.py create mode 100644 tests/test_models.py create mode 100644 tests/test_pack.py create mode 100644 tests/test_worker.py create mode 100644 uv.lock diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b5f08c2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +*.py[cod] +.venv/ +.uv-cache/ +dist/ +build/ +*.zip +.pytest_cache/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..507ed6b --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,25 @@ +## 总规范 + +- 说明性内容（例如提交说明、文档、注释、mermaid），以中文为主，专有名词、约定俗称可用英文 + +## Markdown 规范 + +- 不使用`---`分割器 +- `# Heading` 一级标题仅用于开头 +- mermaid 节点名用英文字符 + +## 代码规范 + +- 避免嵌套结构 +- 特殊的、非常规的需要写注释 +- 变量、函数名、Docstring 用全英文 +- 同一事物用统一表达 +- 模块、函数、变量名，前缀针对不同的命名规范，使用`O *`、`o_*`、`O_*`、`o.*`、`o-*` +- 使用的依赖越少越好 +- 选用简单直接的方式实现 +- 代码、执行方式改变，调整已有的文档、计划，而非新建 + +## Shell 规范 + +- 优先使用、提供 Git Bash 执行所有 Shell 命令 +- Windows 路径在 Bash 命令中优先使用正斜杠形式 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..83c8425 --- /dev/null +++ b/LICENSE @@ -0,0 +1,8 @@ +GNU GENERAL PUBLIC LICENSE +Version 3, 29 June 2007 + +This project is licensed under the GNU General Public License, +version 3 or any later version. + +The complete license text is available at: +https://www.gnu.org/licenses/gpl-3.0.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..55b43ff --- /dev/null +++ b/README.md @@ -0,0 +1,154 @@ +# OmooLab DA3 + +`OmooLab DA3` 是基于 +[Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3) +的 Blender Extension。Extension 在独立 Runtime 中运行 DA3，将官方导出结果保存在 +独立 job 目录，并在任务完成后使用系统文件管理器打开该目录。 + +## 当前功能 + +- 使用 Blender Extension manifest，可通过 `Install from Disk` 安装 +- 在独立 Python Runtime 中运行 Torch 与 DA3，避免污染 Blender Python 进程 +- Runtime、模型和任务结果默认保存在 `~/.omoo_da3`，可由多个 Blender 版本共享 +- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1 +- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入 +- 视频均匀抽帧，可设置参与重建的最大帧数 +- `Setup` 集中管理 Runtime，并逐项列出和下载所有可用模型 +- Runtime 已就绪时不额外占用面板空间；模型列表每个模型只占一行 +- `Debug` 是独立且默认折叠的 DA3 测试面板 +- 使用一个导出类型下拉菜单和一个 `Export` 按钮 +- 只显示当前格式的专属参数；GLB 与 PLY 设置不会同时出现 +- Debug 的模型选项按格式过滤；PLY 只提供两个带 Gaussian head 的 Giant 模型 +- Debug 在模型选项同一行用勾选或警告图标表示下载状态，不显示状态文字、License 和 + 已完成任务的文件计数 +- 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter +- 导出完成后只打开对应 job 文件夹，不向当前 Blender 文件添加任何数据 +- 完整保留每次任务生成的所有官方文件 +- 安装、下载和推理均在后台子进程运行，状态栏显示阶段、百分比进度条与取消按钮 + +Blender 端不读取或转换导出文件，也不会重新反投影深度、构造网格、转换坐标、创建材质或 +Geometry Nodes。 + +## 使用 + +1. 使用 `uv run python pack.py` 构建 Extension ZIP。 +2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 +3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。 +4. 点击 `Install Runtime`。首次安装会下载 Python、PyTorch 和 DA3 依赖。 +5. 在 3D View 按 `N`，打开 `Omoo DA3` 页签。 +6. 在 `Setup` 中找到需要的模型并点击该行的 `Download`。 +7. 展开默认收起的 `Debug`，选择 `Export Format` 和本次执行使用的模型。 +8. 选择一张图片或一个视频文件。视频可设置均匀抽取的帧数，然后点击 `Export`。 + +选择 PLY 后，模型菜单只会提供 `DA3 Giant 1.1` 和 +`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head，无法由 +DA3 生成 `gs_ply`。 + +每次任务会在 `~/.omoo_da3/jobs/<任务目录>/` 中保留 DA3 官方产物： + +```text +/ +└── scene.glb + +/ +└── gs_ply/*.ply + +/ +└── depth_vis/*.jpg + +/ +└── exports/mini_npz/results.npz +``` + +建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重；`DA3 Large 1.1` 权重采用 +CC BY-NC 4.0，仅限非商业用途。 + +## 输出格式说明 + +`Export Format` 下拉菜单分别把以下原始 `export_format` 传给 DA3： + +- GLB：`glb` +- PLY：`gs_ply`，同时传入 `infer_gs=True` +- Depth Images：`depth_vis` +- Mini NPZ：`mini_npz` + +界面使用 DA3 API 的参数名、选项和默认值。GLB 专属设置仅在选择 `GLB` 时显示， +`gs_views_interval` 仅在选择 `PLY (3DGS)` 时显示： + +- `process_res=504` +- `process_res_method="upper_bound_resize"` +- `use_ray_pose=False` +- `ref_view_strategy="saddle_balanced"` +- `conf_thresh_percentile=40.0` +- `num_max_points=1_000_000` +- `show_cameras=True` +- `export_kwargs["gs_ply"]["gs_views_interval"]=1` + +DA3 官方 `glb` 是彩色点云与可选相机线框；`gs_ply` 是带 Gaussian 属性的 PLY； +`depth_vis` 是原图与着色深度图的并排预览；`mini_npz` 保存深度、置信度和可用的相机参数。 +GLB 任务会关闭其默认附带的深度预览，使四种 job 产物保持分离。 + +官方格式与限制可参阅 +[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。 + +## Runtime + +Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Python 3.12 +环境，并安装： + +- PyTorch 2.9.1 +- TorchVision 0.24.1 +- xFormers 0.0.33.post2（与 PyTorch 2.9.1 对应） +- 固定到已验证 commit 的 Depth Anything 3 +- 固定到 DA3 官方推荐 commit 的 `gsplat`；仅安装 PLY 导出所需的纯 Python 部分， + 不编译本机 CUDA rasterizer +- `addict` 2.4.0（补充 DA3 上游包元数据暂未声明的运行依赖） + +默认 CUDA Backend 为 CUDA 12.8，也可以选择 CPU。CPU 可以用于验证，但实际推理 +通常很慢。 + +Runtime、模型和任务结果默认保存在用户主目录的 `~/.omoo_da3`： + +```text +~/.omoo_da3/ +├── runtime/ +├── models/ +└── jobs/ +``` + +不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension +Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录， +因为该目录可能只读，并且升级时会被替换。不要同时从多个 Blender 实例执行 +`Install Runtime`；推理任务和模型读取可以共享。 + +## 开发 + +```bash +uv run python -m unittest discover -s tests -v +uv run python pack.py +``` + +项目结构： + +```text +. +├── pack.py +├── src/ +│ └── omoo_da3/ +│ ├── blender_manifest.toml +│ ├── bootstrap_runtime.py +│ ├── models.py +│ ├── operators.py +│ ├── panel.py +│ ├── preferences.py +│ ├── properties.py +│ ├── runtime.py +│ └── worker.py +└── tests/ +``` + +## 上游与许可证 + +- Depth Anything 3 代码采用 Apache 2.0。 +- 本 Extension 代码采用 GPL-3.0-or-later。 +- 模型权重许可证由所选 Hugging Face 模型仓库决定，Extension 下载前会在界面显示。 diff --git a/pack.py b/pack.py new file mode 100644 index 0000000..a0f76b8 --- /dev/null +++ b/pack.py @@ -0,0 +1,32 @@ +import tomllib +from pathlib import Path +from zipfile import ZIP_DEFLATED, ZipFile + + +PROJECT_ROOT = Path(__file__).resolve().parent +SOURCE_DIR = PROJECT_ROOT / "src" / "omoo_da3" +DIST_DIR = PROJECT_ROOT / "dist" + + +def build_extension(): + manifest = tomllib.loads( + (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8") + ) + output_path = DIST_DIR / f"{manifest['id']}-{manifest['version']}.zip" + DIST_DIR.mkdir(exist_ok=True) + + with ZipFile(output_path, "w", ZIP_DEFLATED) as archive: + for source_path in sorted(SOURCE_DIR.rglob("*")): + if not source_path.is_file(): + continue + if "__pycache__" in source_path.parts: + continue + archive.write(source_path, source_path.relative_to(SOURCE_DIR)) + + return output_path + + +if __name__ == "__main__": + extension_path = build_extension() + size_mib = extension_path.stat().st_size / (1024 * 1024) + print(f"Built {extension_path} ({size_mib:.2f} MiB)") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..32e4728 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "omoo-da3" +version = "0.8.1" +description = "Direct Depth Anything 3 exports from Blender" +readme = "README.md" +requires-python = ">=3.11" +license = "GPL-3.0-or-later" +dependencies = [] + +[tool.uv] +package = false diff --git a/src/omoo_da3/__init__.py b/src/omoo_da3/__init__.py new file mode 100644 index 0000000..4d555c4 --- /dev/null +++ b/src/omoo_da3/__init__.py @@ -0,0 +1,41 @@ +import bpy + +from .operators import ( + OMOO_DA3_OT_cancel_job, + OMOO_DA3_OT_run_job, + cancel_active_job, + close_active_job, +) +from .panel import OMOO_DA3_PT_debug, OMOO_DA3_PT_runtime_models +from .preferences import OMOO_DA3_preferences +from .properties import OMOO_DA3_PG_settings +from .status_bar import draw_status_bar + + +CLASSES = ( + OMOO_DA3_preferences, + OMOO_DA3_PG_settings, + OMOO_DA3_OT_run_job, + OMOO_DA3_OT_cancel_job, + OMOO_DA3_PT_runtime_models, + OMOO_DA3_PT_debug, +) + + +def register(): + for class_type in CLASSES: + bpy.utils.register_class(class_type) + bpy.types.Scene.omoo_da3_settings = bpy.props.PointerProperty( + type=OMOO_DA3_PG_settings, + ) + bpy.types.STATUSBAR_HT_header.prepend(draw_status_bar) + + +def unregister(): + cancel_active_job() + close_active_job() + bpy.types.STATUSBAR_HT_header.remove(draw_status_bar) + if hasattr(bpy.types.Scene, "omoo_da3_settings"): + del bpy.types.Scene.omoo_da3_settings + for class_type in reversed(CLASSES): + bpy.utils.unregister_class(class_type) diff --git a/src/omoo_da3/blender_manifest.toml b/src/omoo_da3/blender_manifest.toml new file mode 100644 index 0000000..2b8b519 --- /dev/null +++ b/src/omoo_da3/blender_manifest.toml @@ -0,0 +1,23 @@ +schema_version = "1.0.0" + +id = "omoo_da3" +version = "0.8.1" +name = "OmooLab DA3" +tagline = "Export official DA3 GLB, PLY, depth, and mini NPZ" +maintainer = "OmooLab" +type = "add-on" + +blender_version_min = "4.5.0" +platforms = ["windows-x64"] +license = ["SPDX:GPL-3.0-or-later"] + +[permissions] +network = "Download DA3 runtime dependencies and model weights" +files = "Read source images and store runtime data" + +[build] +paths_exclude_pattern = [ + "__pycache__/", + ".*", + "*.zip", +] diff --git a/src/omoo_da3/bootstrap_runtime.py b/src/omoo_da3/bootstrap_runtime.py new file mode 100644 index 0000000..6187e76 --- /dev/null +++ b/src/omoo_da3/bootstrap_runtime.py @@ -0,0 +1,203 @@ +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + + +TORCH_VERSION = "2.9.1" +TORCHVISION_VERSION = "0.24.1" +XFORMERS_VERSION = "0.0.33.post2" +ADDICT_VERSION = "2.4.0" +GSPLAT_COMMIT = "0b4dddf04cb687367602c01196913cde6a743d70" +RUNTIME_SCHEMA = 2 +TORCH_INDEXES = { + "cu128": "https://download.pytorch.org/whl/cu128", + "cpu": "https://download.pytorch.org/whl/cpu", +} + + +def write_status(path, progress, message, error=None): + payload = {"progress": progress, "message": message} + if error: + payload["error"] = error + Path(path).write_text( + json.dumps(payload, ensure_ascii=False), + encoding="utf-8", + ) + + +def run(command, environment=None): + print("Running:", subprocess.list2cmdline([str(value) for value in command])) + process_environment = None + if environment: + process_environment = os.environ.copy() + process_environment.update(environment) + subprocess.check_call( + [str(value) for value in command], + env=process_environment, + ) + + +def environment_python(runtime_dir): + if sys.platform == "win32": + return Path(runtime_dir) / "Scripts" / "python.exe" + return Path(runtime_dir) / "bin" / "python" + + +def install(args): + runtime_dir = Path(args.runtime_dir) + runtime_dir.parent.mkdir(parents=True, exist_ok=True) + + write_status(args.status, 0.03, "Creating isolated Python Runtime") + run( + [ + args.uv, + "venv", + "--clear", + "--python", + "3.12", + runtime_dir, + ] + ) + python = environment_python(runtime_dir) + + write_status(args.status, 0.12, "Installing PyTorch") + run( + [ + args.uv, + "pip", + "install", + "--python", + python, + f"torch=={TORCH_VERSION}", + f"torchvision=={TORCHVISION_VERSION}", + "--index-url", + TORCH_INDEXES[args.backend], + ] + ) + + write_status(args.status, 0.38, "Installing matching xFormers") + xformers_command = [ + args.uv, + "pip", + "install", + "--python", + python, + f"xformers=={XFORMERS_VERSION}", + ] + if args.backend == "cu128": + xformers_command.extend( + ["--index-url", TORCH_INDEXES[args.backend]] + ) + else: + xformers_command.append("--no-deps") + run(xformers_command) + + write_status(args.status, 0.50, "Installing Depth Anything 3") + da3_requirement = ( + "depth-anything-3 @ " + "git+https://github.com/ByteDance-Seed/Depth-Anything-3.git" + f"@{args.da3_commit}" + ) + run( + [ + args.uv, + "pip", + "install", + "--python", + python, + da3_requirement, + f"addict=={ADDICT_VERSION}", + ] + ) + + write_status(args.status, 0.78, "Installing gsplat build tools") + run( + [ + args.uv, + "pip", + "install", + "--python", + python, + "setuptools", + "wheel", + "ninja", + ] + ) + + write_status(args.status, 0.84, "Installing gsplat") + # PLY 导出不调用 rasterizer，禁用本地扩展可避免 CUDA 与 MSVC 工具链耦合。 + gsplat_environment = {"BUILD_NO_CUDA": "1"} + run( + [ + args.uv, + "pip", + "install", + "--python", + python, + "--no-build-isolation", + ( + "gsplat @ git+https://github.com/nerfstudio-project/" + f"gsplat.git@{GSPLAT_COMMIT}" + ), + ], + gsplat_environment, + ) + + write_status(args.status, 0.92, "Checking Runtime") + run( + [ + python, + "-c", + ( + "import torch; " + "from depth_anything_3.api import DepthAnything3; " + "import gsplat; " + "print(torch.__version__, torch.cuda.is_available())" + ), + ] + ) + + ready_path = runtime_dir / "omoo_da3_ready.json" + ready_path.write_text( + json.dumps( + { + "runtime_schema": RUNTIME_SCHEMA, + "python": "3.12", + "torch": TORCH_VERSION, + "xformers": XFORMERS_VERSION, + "backend": args.backend, + "da3_commit": args.da3_commit, + "gsplat_commit": GSPLAT_COMMIT, + "gsplat_cuda": False, + }, + indent=2, + ), + encoding="utf-8", + ) + write_status(args.status, 1.0, "Runtime is ready") + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--uv", required=True) + parser.add_argument("--runtime-dir", required=True) + parser.add_argument("--backend", choices=tuple(TORCH_INDEXES), required=True) + parser.add_argument("--da3-commit", required=True) + parser.add_argument("--status", required=True) + return parser.parse_args() + + +def main(): + args = parse_args() + try: + install(args) + except Exception as error: + write_status(args.status, 1.0, "Runtime installation failed", str(error)) + raise + + +if __name__ == "__main__": + main() diff --git a/src/omoo_da3/models.py b/src/omoo_da3/models.py new file mode 100644 index 0000000..8e86095 --- /dev/null +++ b/src/omoo_da3/models.py @@ -0,0 +1,81 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ModelSpec: + key: str + label: str + repository: str + license_name: str + description: str + supports_gs: bool + + +MODELS = { + "SMALL": ModelSpec( + key="SMALL", + label="DA3 Small", + repository="depth-anything/DA3-SMALL", + license_name="Apache 2.0", + description="Fastest any-view model with Apache 2.0 weights", + supports_gs=False, + ), + "BASE": ModelSpec( + key="BASE", + label="DA3 Base", + repository="depth-anything/DA3-BASE", + license_name="Apache 2.0", + description="Balanced any-view model with Apache 2.0 weights", + supports_gs=False, + ), + "LARGE_1_1": ModelSpec( + key="LARGE_1_1", + label="DA3 Large 1.1", + repository="depth-anything/DA3-LARGE-1.1", + license_name="CC BY-NC 4.0", + description="Higher quality non-commercial model with larger VRAM use", + supports_gs=False, + ), + "GIANT_1_1": ModelSpec( + key="GIANT_1_1", + label="DA3 Giant 1.1", + repository="depth-anything/DA3-GIANT-1.1", + license_name="CC BY-NC 4.0", + description="Non-commercial Giant model with Gaussian head", + supports_gs=True, + ), + "NESTED_GIANT_LARGE_1_1": ModelSpec( + key="NESTED_GIANT_LARGE_1_1", + label="DA3 Nested Giant Large 1.1", + repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1", + license_name="CC BY-NC 4.0", + description="Non-commercial nested model with metric depth and Gaussian head", + supports_gs=True, + ), +} + + +def enum_items(supports_gs=None): + return [ + ( + model.key, + model.label, + f"{model.description}; license: {model.license_name}", + ) + for model in MODELS.values() + if supports_gs is None or model.supports_gs == supports_gs + ] + + +def selected_model_key(settings, export_format=None): + selected_format = export_format or settings.export_format + if selected_format == "gs_ply": + return settings.gs_model + return settings.model + + +def get_model(model_key): + try: + return MODELS[model_key] + except KeyError as error: + raise ValueError(f"Unknown DA3 model: {model_key}") from error diff --git a/src/omoo_da3/operators.py b/src/omoo_da3/operators.py new file mode 100644 index 0000000..3f5fe6d --- /dev/null +++ b/src/omoo_da3/operators.py @@ -0,0 +1,283 @@ +import subprocess +from dataclasses import dataclass +from pathlib import Path + +import bpy + +from . import runtime +from .models import get_model, selected_model_key + + +RECONSTRUCTION_JOB = "RECONSTRUCT" + + +def open_job_directory(result_directory): + result_directory = Path(result_directory) + if not result_directory.is_dir(): + raise RuntimeError(f"Job directory does not exist: {result_directory}") + result = bpy.ops.wm.path_open(filepath=str(result_directory)) + if "FINISHED" not in result: + raise RuntimeError(f"Unable to open job directory: {result_directory}") + + +@dataclass +class ActiveJob: + kind: str + process: subprocess.Popen + log_file: object + log_path: Path + status_path: Path + result_directory: Path + export_format: str = "" + cancelled: bool = False + + +_active_job = None + + +def active_job(): + return _active_job + + +def cancel_active_job(): + global _active_job + + if _active_job is None: + return + _active_job.cancelled = True + if _active_job.process.poll() is None: + _active_job.process.terminate() + + +def close_active_job(): + global _active_job + + if _active_job is None: + return + _active_job.log_file.close() + _active_job = None + + +class OMOO_DA3_OT_run_job(bpy.types.Operator): + bl_idname = "omoo_da3.run_job" + bl_label = "Run Omoo DA3 Job" + bl_options = {"INTERNAL"} + + job_type: bpy.props.EnumProperty( + items=( + ("INSTALL_RUNTIME", "Install Runtime", ""), + ("DOWNLOAD_MODEL", "Download Model", ""), + (RECONSTRUCTION_JOB, "Export DA3 Result", ""), + ), + options={"HIDDEN", "SKIP_SAVE"}, + ) + model_key: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + + def execute(self, context): + global _active_job + + if active_job() is not None: + self.report({"WARNING"}, "Another Omoo DA3 task is running") + return {"CANCELLED"} + + settings = context.scene.omoo_da3_settings + try: + command, paths = self._prepare_command(context, settings) + except (OSError, RuntimeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + + log_file = paths["log"].open("w", encoding="utf-8") + try: + process = subprocess.Popen( + command, + stdout=log_file, + stderr=subprocess.STDOUT, + cwd=runtime.storage_root(), + creationflags=runtime.creation_flags(), + ) + except OSError as error: + log_file.close() + self.report({"ERROR"}, f"Unable to start DA3 task: {error}") + return {"CANCELLED"} + + _active_job = ActiveJob( + kind=self.job_type, + process=process, + log_file=log_file, + log_path=paths["log"], + status_path=paths["status"], + result_directory=paths["result"], + export_format=( + settings.export_format + if self.job_type == RECONSTRUCTION_JOB + else "" + ), + ) + settings.status = self._starting_message(settings) + settings.progress = 0.0 + self._timer = context.window_manager.event_timer_add( + 0.25, + window=context.window, + ) + context.window_manager.modal_handler_add(self) + self._redraw_ui(context, force=True) + return {"RUNNING_MODAL"} + + def _prepare_command(self, _context, settings): + paths = runtime.new_job_paths(self.job_type) + if self.job_type == "INSTALL_RUNTIME": + self._require_online_access() + return runtime.install_runtime_command(paths["status"]), paths + + if not runtime.runtime_ready(): + raise RuntimeError("Install the Omoo DA3 Runtime first") + + if self.job_type == "DOWNLOAD_MODEL": + self._require_online_access() + model = get_model(self.model_key) + return ( + runtime.download_model_command(model.key, paths["status"]), + paths, + ) + + input_path = Path(bpy.path.abspath(settings.input_path)) + if not input_path.is_file(): + raise ValueError("Select a valid image or video file") + model_key = selected_model_key(settings) + model = get_model(model_key) + if not runtime.model_ready(model_key): + raise RuntimeError(f"Download {model.label} first") + if self.job_type != RECONSTRUCTION_JOB: + raise RuntimeError(f"Unsupported task type: {self.job_type}") + export_format = settings.export_format + if export_format == "gs_ply" and not model.supports_gs: + raise RuntimeError( + "PLY export requires a DA3 Giant or Nested Giant model" + ) + return ( + runtime.reconstruct_command( + settings, + export_format, + paths["status"], + paths["result"], + ), + paths, + ) + + def _require_online_access(self): + if not bpy.app.online_access: + raise RuntimeError( + "Online Access is disabled in Blender Preferences" + ) + + def _starting_message(self, settings): + labels = { + "INSTALL_RUNTIME": "Starting Runtime installation", + "DOWNLOAD_MODEL": "Starting model download", + } + if self.job_type == RECONSTRUCTION_JOB: + return f"Starting DA3 {settings.export_format} export" + if self.job_type == "DOWNLOAD_MODEL": + return f"Starting {get_model(self.model_key).label} download" + return labels[self.job_type] + + def modal(self, context, event): + if event.type == "ESC": + cancel_active_job() + + if event.type != "TIMER": + return {"PASS_THROUGH"} + + job = active_job() + if job is None: + self._remove_timer(context) + return {"CANCELLED"} + + self._update_status(context, job) + return_code = job.process.poll() + if return_code is None: + return {"RUNNING_MODAL"} + return self._finish(context, job, return_code) + + def _update_status(self, context, job): + status = runtime.read_status(job.status_path) + if status is None: + return + settings = context.scene.omoo_da3_settings + settings.progress = float(status.get("progress", settings.progress)) + settings.status = status.get("message", settings.status) + self._redraw_ui(context) + + def _finish(self, context, job, return_code): + settings = context.scene.omoo_da3_settings + self._remove_timer(context) + job.log_file.close() + + if job.cancelled: + settings.status = "Task cancelled" + settings.progress = 0.0 + close_active_job() + self._redraw_ui(context, force=True) + return {"CANCELLED"} + + status = runtime.read_status(job.status_path) or {} + error = status.get("error") + if return_code != 0 or error: + settings.status = "Task failed" + settings.progress = 0.0 + self.report( + {"ERROR"}, + f"{error or 'DA3 subprocess failed'}; log: {job.log_path}", + ) + close_active_job() + self._redraw_ui(context, force=True) + return {"CANCELLED"} + + if job.kind == RECONSTRUCTION_JOB: + try: + open_job_directory(job.result_directory) + except Exception as open_error: + self.report({"WARNING"}, str(open_error)) + self.report( + {"INFO"}, + f"Exported {job.export_format} to {job.result_directory}", + ) + + settings.status = status.get("message", "Task complete") + settings.progress = 1.0 + close_active_job() + self._redraw_ui(context, force=True) + return {"FINISHED"} + + def _remove_timer(self, context): + timer = getattr(self, "_timer", None) + if timer is None: + return + context.window_manager.event_timer_remove(timer) + self._timer = None + + def _redraw_ui(self, context, force=False): + from .status_bar import redraw_ui + + redraw_ui(context, force=force) + + def cancel(self, context): + cancel_active_job() + self._remove_timer(context) + self._redraw_ui(context, force=True) + + +class OMOO_DA3_OT_cancel_job(bpy.types.Operator): + bl_idname = "omoo_da3.cancel_job" + bl_label = "Cancel Omoo DA3 Task" + + @classmethod + def poll(cls, _context): + return active_job() is not None + + def execute(self, _context): + cancel_active_job() + return {"FINISHED"} diff --git a/src/omoo_da3/panel.py b/src/omoo_da3/panel.py new file mode 100644 index 0000000..03b80cb --- /dev/null +++ b/src/omoo_da3/panel.py @@ -0,0 +1,112 @@ +import bpy + +from .models import MODELS, selected_model_key +from .operators import active_job +from .runtime import model_ready, runtime_ready + + +class OMOO_DA3_PT_runtime_models(bpy.types.Panel): + bl_label = "Setup" + bl_idname = "OMOO_DA3_PT_runtime_models" + bl_space_type = "VIEW_3D" + bl_region_type = "UI" + bl_category = "Omoo DA3" + + def draw(self, _context): + layout = self.layout + running = active_job() is not None + + if not runtime_ready(): + layout.label(text="Runtime is not installed", icon="ERROR") + if not running: + operator = layout.operator( + "omoo_da3.run_job", + text="Install Runtime", + icon="IMPORT", + ) + operator.job_type = "INSTALL_RUNTIME" + + for model in MODELS.values(): + self._draw_model(layout, model, running) + + def _draw_model(self, layout, model, running): + row = layout.row(align=True) + row.label(text=model.label) + if model_ready(model.key): + row.label(text="Ready", icon="CHECKMARK") + elif runtime_ready() and not running: + operator = row.operator( + "omoo_da3.run_job", + text="Download", + icon="IMPORT", + ) + operator.job_type = "DOWNLOAD_MODEL" + operator.model_key = model.key + else: + row.label(text="Not downloaded", icon="ERROR") + + +class OMOO_DA3_PT_debug(bpy.types.Panel): + bl_label = "Debug" + bl_idname = "OMOO_DA3_PT_debug" + bl_space_type = "VIEW_3D" + bl_region_type = "UI" + bl_category = "Omoo DA3" + bl_options = {"DEFAULT_CLOSED"} + + def draw(self, context): + layout = self.layout + settings = context.scene.omoo_da3_settings + + layout.prop(settings, "export_format") + model_key = selected_model_key(settings) + model_property = ( + "gs_model" if settings.export_format == "gs_ply" else "model" + ) + model_row = layout.row(align=True) + model_row.prop(settings, model_property) + if model_ready(model_key): + model_row.label(text="", icon="CHECKMARK") + else: + model_row.label(text="", icon="ERROR") + + layout.prop(settings, "input_path") + input_path = bpy.path.abspath(settings.input_path) + if input_path.lower().endswith( + ( + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", + ) + ): + layout.prop(settings, "video_frames") + layout.prop(settings, "device") + layout.prop(settings, "process_res") + layout.prop(settings, "process_res_method") + layout.prop(settings, "use_ray_pose") + layout.prop(settings, "ref_view_strategy") + + if settings.export_format == "glb": + glb_box = layout.box() + glb_box.label(text="GLB Settings") + glb_box.prop(settings, "conf_thresh_percentile") + glb_box.prop(settings, "num_max_points") + glb_box.prop(settings, "show_cameras") + elif settings.export_format == "gs_ply": + ply_box = layout.box() + ply_box.label(text="PLY (3DGS) Settings") + ply_box.prop(settings, "gs_views_interval") + + if active_job() is None: + operator = layout.operator( + "omoo_da3.run_job", + text="Export", + icon="EXPORT", + ) + operator.job_type = "RECONSTRUCT" diff --git a/src/omoo_da3/preferences.py b/src/omoo_da3/preferences.py new file mode 100644 index 0000000..4c70f4f --- /dev/null +++ b/src/omoo_da3/preferences.py @@ -0,0 +1,52 @@ +import shutil +from pathlib import Path + +import bpy + + +class OMOO_DA3_preferences(bpy.types.AddonPreferences): + bl_idname = __package__ + + storage_root: bpy.props.StringProperty( + name="Storage Root", + description=( + "Shared user-writable storage for Runtime, models and jobs; " + "all Blender versions can use the same directory" + ), + subtype="DIR_PATH", + default=str(Path.home() / ".omoo_da3"), + ) + uv_executable: bpy.props.StringProperty( + name="uv Executable", + description="Path to uv; a command name is accepted when uv is on PATH", + subtype="FILE_PATH", + default=shutil.which("uv") or "uv", + ) + torch_backend: bpy.props.EnumProperty( + name="Torch Backend", + description="PyTorch build installed into the isolated Runtime", + items=( + ("CU128", "CUDA 12.8", "NVIDIA CUDA 12.8 PyTorch wheels"), + ("CPU", "CPU", "CPU-only PyTorch wheels"), + ), + default="CU128", + ) + + def draw(self, _context): + from .runtime import runtime_directory, runtime_ready + + layout = self.layout + layout.prop(self, "storage_root") + layout.prop(self, "uv_executable") + layout.prop(self, "torch_backend") + layout.separator() + layout.label(text=f"Runtime: {runtime_directory()}") + if runtime_ready(): + layout.label(text="Runtime is ready", icon="CHECKMARK") + else: + operator = layout.operator( + "omoo_da3.run_job", + text="Install Runtime", + icon="IMPORT", + ) + operator.job_type = "INSTALL_RUNTIME" diff --git a/src/omoo_da3/properties.py b/src/omoo_da3/properties.py new file mode 100644 index 0000000..f0028bb --- /dev/null +++ b/src/omoo_da3/properties.py @@ -0,0 +1,146 @@ +import bpy + +from .models import enum_items + + +class OMOO_DA3_PG_settings(bpy.types.PropertyGroup): + input_path: bpy.props.StringProperty( + name="Image or Video", + description="A single image or video file used for reconstruction", + subtype="FILE_PATH", + ) + video_frames: bpy.props.IntProperty( + name="Video Frames", + description="Maximum number of frames sampled evenly from a video", + default=16, + min=2, + soft_max=64, + ) + model: bpy.props.EnumProperty( + name="Model", + description="Depth Anything 3 model", + items=enum_items(), + default="BASE", + ) + gs_model: bpy.props.EnumProperty( + name="Model", + description="Depth Anything 3 model with a Gaussian head", + items=enum_items(supports_gs=True), + default="GIANT_1_1", + ) + device: bpy.props.EnumProperty( + name="Device", + description="Inference device used by the external Runtime", + items=( + ("AUTO", "Auto", "Use CUDA when available, otherwise CPU"), + ("CUDA", "CUDA", "Require an NVIDIA CUDA device"), + ("CPU", "CPU", "Run on CPU"), + ), + default="AUTO", + ) + process_res: bpy.props.IntProperty( + name="Process Res", + description="DA3 base processing resolution", + default=504, + min=1, + soft_max=1008, + ) + process_res_method: bpy.props.EnumProperty( + name="Process Res Method", + description="DA3 image resize method", + items=( + ( + "upper_bound_resize", + "Upper Bound Resize", + "Make Process Res the longer side", + ), + ( + "lower_bound_resize", + "Lower Bound Resize", + "Make Process Res the shorter side", + ), + ), + default="upper_bound_resize", + ) + conf_thresh_percentile: bpy.props.FloatProperty( + name="Conf Thresh Percentile", + description="DA3 GLB confidence threshold percentile", + default=40.0, + min=0.0, + max=100.0, + subtype="PERCENTAGE", + ) + num_max_points: bpy.props.IntProperty( + name="Num Max Points", + description="DA3 GLB maximum point count", + default=1_000_000, + min=1, + soft_max=5_000_000, + ) + show_cameras: bpy.props.BoolProperty( + name="Show Cameras", + description="Include DA3 camera wireframes in the GLB export", + default=True, + ) + use_ray_pose: bpy.props.BoolProperty( + name="Use Ray Pose", + description="Use the slower and more accurate ray-based pose estimation", + default=False, + ) + ref_view_strategy: bpy.props.EnumProperty( + name="Ref View Strategy", + description="DA3 reference view selection strategy", + items=( + ("first", "First", "Always use the first view"), + ("middle", "Middle", "Use the middle view"), + ( + "saddle_balanced", + "Saddle Balanced", + "Select a view with balanced feature metrics", + ), + ( + "saddle_sim_range", + "Saddle Sim Range", + "Select the view with the largest similarity range", + ), + ), + default="saddle_balanced", + ) + export_format: bpy.props.EnumProperty( + name="Export Format", + description="Official DA3 export format", + items=( + ("glb", "GLB", "Export a colored point cloud as GLB"), + ("gs_ply", "PLY (3DGS)", "Export Gaussian splats as PLY"), + ( + "depth_vis", + "Depth Images", + "Export image and depth visualization pairs", + ), + ( + "mini_npz", + "Mini NPZ", + "Export compact depth and camera data", + ), + ), + default="glb", + ) + gs_views_interval: bpy.props.IntProperty( + name="GS Views Interval", + description="Export 3D Gaussians every N input views", + default=1, + min=1, + ) + status: bpy.props.StringProperty( + name="Status", + default="Ready", + options={"SKIP_SAVE"}, + ) + progress: bpy.props.FloatProperty( + name="Progress", + default=0.0, + min=0.0, + max=1.0, + subtype="FACTOR", + options={"SKIP_SAVE"}, + ) diff --git a/src/omoo_da3/runtime.py b/src/omoo_da3/runtime.py new file mode 100644 index 0000000..3716ba9 --- /dev/null +++ b/src/omoo_da3/runtime.py @@ -0,0 +1,198 @@ +import json +import os +import shutil +import subprocess +import uuid +from datetime import datetime +from pathlib import Path + +import bpy + +from .models import get_model, selected_model_key + + +DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" +RUNTIME_SCHEMA = 2 +DEFAULT_STORAGE_ROOT = Path.home() / ".omoo_da3" + + +def preferences(): + addon = bpy.context.preferences.addons.get(__package__) + return addon.preferences if addon is not None else None + + +def storage_root(): + addon_preferences = preferences() + configured = getattr(addon_preferences, "storage_root", "") + root = ( + Path(bpy.path.abspath(configured)).expanduser() + if configured + else DEFAULT_STORAGE_ROOT + ) + root.mkdir(parents=True, exist_ok=True) + return root + + +def runtime_directory(): + return storage_root() / "runtime" + + +def runtime_python(): + directory = runtime_directory() + if os.name == "nt": + return directory / "Scripts" / "python.exe" + return directory / "bin" / "python" + + +def runtime_ready(): + ready_path = runtime_directory() / "omoo_da3_ready.json" + if not runtime_python().is_file() or not ready_path.is_file(): + return False + try: + ready = json.loads(ready_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return False + return ready.get("runtime_schema") == RUNTIME_SCHEMA + + +def models_directory(): + path = storage_root() / "models" + path.mkdir(parents=True, exist_ok=True) + return path + + +def model_directory(model_key): + return models_directory() / model_key.lower() + + +def model_ready(model_key): + path = model_directory(model_key) + return (path / "config.json").is_file() and any( + path.glob("*.safetensors") + ) + + +def jobs_directory(): + path = storage_root() / "jobs" + path.mkdir(parents=True, exist_ok=True) + return path + + +def new_job_paths(job_type): + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + token = uuid.uuid4().hex[:8] + stem = f"{timestamp}-{job_type.lower()}-{token}" + directory = jobs_directory() + return { + "log": directory / f"{stem}.log", + "status": directory / f"{stem}.json", + "result": directory / stem, + } + + +def uv_executable(): + addon_preferences = preferences() + configured = getattr(addon_preferences, "uv_executable", "uv") + resolved = shutil.which(configured) + if resolved: + return resolved + if Path(configured).is_file(): + return str(Path(configured)) + raise RuntimeError("uv executable was not found") + + +def install_runtime_command(status_path): + executable = uv_executable() + addon_preferences = preferences() + backend = getattr(addon_preferences, "torch_backend", "CU128") + script = Path(__file__).with_name("bootstrap_runtime.py") + return [ + executable, + "run", + "--no-project", + "--python", + "3.12", + str(script), + "--uv", + executable, + "--runtime-dir", + str(runtime_directory()), + "--backend", + backend.lower(), + "--da3-commit", + DA3_COMMIT, + "--status", + str(status_path), + ] + + +def download_model_command(model_key, status_path): + model = get_model(model_key) + return [ + str(runtime_python()), + str(Path(__file__).with_name("worker.py")), + "download", + "--repository", + model.repository, + "--model-dir", + str(model_directory(model_key)), + "--status", + str(status_path), + ] + + +def reconstruct_command( + settings, + export_format, + status_path, + result_directory, +): + model_key = selected_model_key(settings, export_format) + return [ + str(runtime_python()), + str(Path(__file__).with_name("worker.py")), + "reconstruct", + "--input", + str(Path(bpy.path.abspath(settings.input_path))), + "--model-dir", + str(model_directory(model_key)), + "--output-dir", + str(result_directory), + "--export-format", + export_format, + "--status", + str(status_path), + "--process-res", + str(settings.process_res), + "--process-res-method", + settings.process_res_method, + "--conf-thresh-percentile", + str(settings.conf_thresh_percentile), + "--num-max-points", + str(settings.num_max_points), + "--show-cameras", + "1" if settings.show_cameras else "0", + "--gs-views-interval", + str(settings.gs_views_interval), + "--video-frames", + str(settings.video_frames), + "--device", + settings.device.lower(), + "--use-ray-pose", + "1" if settings.use_ray_pose else "0", + "--ref-view-strategy", + settings.ref_view_strategy, + ] + + +def creation_flags(): + if os.name != "nt": + return 0 + return subprocess.CREATE_NO_WINDOW + + +def read_status(path): + try: + return json.loads(Path(path).read_text(encoding="utf-8")) + except (FileNotFoundError, json.JSONDecodeError, OSError): + return None diff --git a/src/omoo_da3/status_bar.py b/src/omoo_da3/status_bar.py new file mode 100644 index 0000000..1989b2a --- /dev/null +++ b/src/omoo_da3/status_bar.py @@ -0,0 +1,51 @@ +import bpy + +from .operators import active_job + + +def draw_status_bar(self, context): + scene = getattr(context, "scene", None) + if active_job() is None or scene is None: + return + + settings = getattr(scene, "omoo_da3_settings", None) + if settings is None: + return + + progress = min(max(settings.progress, 0.0), 1.0) + row = self.layout.row(align=True) + row.label(text=f"Omoo DA3: {settings.status}", icon="TIME") + + progress_row = row.row(align=True) + progress_row.ui_units_x = 9.0 + progress_row.progress( + factor=progress, + type="BAR", + text=f"{progress:.0%}", + ) + row.operator( + "omoo_da3.cancel_job", + text="", + icon="X", + ) + + +def redraw_ui(context, force=False): + window_manager = getattr(context, "window_manager", None) + if window_manager is None: + return + + for window in window_manager.windows: + screen = getattr(window, "screen", None) + if screen is None: + continue + for area in screen.areas: + if area.type in {"STATUSBAR", "VIEW_3D"}: + area.tag_redraw() + + if not force or getattr(context, "window", None) is None: + return + try: + bpy.ops.wm.redraw_timer(type="DRAW_WIN_SWAP", iterations=1) + except RuntimeError: + pass diff --git a/src/omoo_da3/worker.py b/src/omoo_da3/worker.py new file mode 100644 index 0000000..aa08759 --- /dev/null +++ b/src/omoo_da3/worker.py @@ -0,0 +1,304 @@ +import argparse +import json +import shutil +import tempfile +import time +from pathlib import Path + + +IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} +VIDEO_SUFFIXES = { + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", +} +SCENE_FILENAME = "scene.glb" +EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz") + + +def write_status(path, progress, message, error=None): + payload = {"progress": progress, "message": message} + if error: + payload["error"] = error + Path(path).write_text( + json.dumps(payload, ensure_ascii=False), + encoding="utf-8", + ) + + +def collect_images(directory): + directory = Path(directory) + return sorted( + path + for path in directory.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + + +def extract_video_frames(video_path, output_directory, maximum_frames, cv2): + capture = cv2.VideoCapture(str(video_path)) + if not capture.isOpened(): + raise RuntimeError(f"Unable to open video: {video_path}") + + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) + if total_frames <= 0: + capture.release() + raise RuntimeError("The video does not report a readable frame count") + + import numpy as np + + sample_count = min(maximum_frames, total_frames) + frame_indices = np.linspace( + 0, + total_frames - 1, + sample_count, + dtype=np.int64, + ) + image_paths = [] + try: + for output_index, frame_index in enumerate(frame_indices): + capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index)) + success, frame = capture.read() + if not success: + raise RuntimeError( + f"Unable to read video frame {int(frame_index)}" + ) + image_path = output_directory / f"frame-{output_index:04d}.jpg" + if not cv2.imwrite(str(image_path), frame): + raise RuntimeError(f"Unable to write video frame: {image_path}") + image_paths.append(image_path) + finally: + capture.release() + + return image_paths + + +def prepare_input(input_path, video_frames, temporary_root): + input_path = Path(input_path) + suffix = input_path.suffix.lower() + if suffix in IMAGE_SUFFIXES: + return [input_path], None + if suffix not in VIDEO_SUFFIXES: + supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES)) + raise RuntimeError(f"Unsupported input type; expected one of: {supported}") + + import cv2 + + temporary_directory = Path( + tempfile.mkdtemp(prefix="omoo-da3-video-", dir=temporary_root) + ) + try: + images = extract_video_frames( + input_path, + temporary_directory, + video_frames, + cv2, + ) + except Exception: + shutil.rmtree(temporary_directory, ignore_errors=True) + raise + return images, temporary_directory + + +def download_model(args): + from huggingface_hub import snapshot_download + + model_dir = Path(args.model_dir) + model_dir.mkdir(parents=True, exist_ok=True) + write_status(args.status, 0.05, f"Downloading {args.repository}") + snapshot_download( + repo_id=args.repository, + local_dir=model_dir, + ) + write_status(args.status, 1.0, "Model is ready") + + +def select_device(requested, torch): + if requested == "cpu": + return torch.device("cpu") + if requested == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available") + return torch.device("cuda") + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def exported_paths(output_directory, export_format): + output_directory = Path(output_directory) + if export_format == "glb": + scene_path = output_directory / SCENE_FILENAME + return [scene_path] if scene_path.is_file() else [] + if export_format == "gs_ply": + return sorted((output_directory / "gs_ply").glob("*.ply")) + if export_format == "depth_vis": + return sorted((output_directory / "depth_vis").glob("*.jpg")) + if export_format == "mini_npz": + result_path = output_directory / "exports" / "mini_npz" / "results.npz" + return [result_path] if result_path.is_file() else [] + return [] + + +def wait_for_exported_paths(output_directory, export_format, timeout=30.0): + # 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。 + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + result_paths = exported_paths(output_directory, export_format) + if result_paths: + return result_paths + time.sleep(0.05) + return [] + + +def reconstruct(args): + import torch + from depth_anything_3.api import DepthAnything3 + from depth_anything_3.utils.export import export + + output_directory = Path(args.output_dir) + output_directory.mkdir(parents=True, exist_ok=False) + write_status(args.status, 0.03, "Preparing input") + images, temporary_directory = prepare_input( + args.input, + args.video_frames, + output_directory.parent, + ) + + try: + device = select_device(args.device, torch) + write_status(args.status, 0.10, f"Loading model on {device.type}") + model = DepthAnything3.from_pretrained(args.model_dir) + model = model.to(device=device) + model.eval() + + write_status(args.status, 0.18, f"Processing {len(images)} frame(s)") + export_kwargs = {} + if args.export_format == "glb": + export_kwargs["glb"] = { + "conf_thresh_percentile": args.conf_thresh_percentile, + "num_max_points": args.num_max_points, + "show_cameras": bool(args.show_cameras), + "export_depth_vis": False, + } + if args.export_format == "gs_ply": + export_kwargs["gs_ply"] = { + "gs_views_interval": args.gs_views_interval, + } + with torch.inference_mode(): + prediction = model.inference( + [str(path) for path in images], + infer_gs=args.export_format == "gs_ply", + process_res=args.process_res, + process_res_method=args.process_res_method, + use_ray_pose=bool(args.use_ray_pose), + ref_view_strategy=args.ref_view_strategy, + ) + write_status(args.status, 0.88, f"Exporting {args.export_format}") + export( + prediction, + args.export_format, + str(output_directory), + **export_kwargs, + ) + + result_paths = wait_for_exported_paths( + output_directory, + args.export_format, + ) + if not result_paths: + raise RuntimeError( + f"DA3 did not export any {args.export_format} files" + ) + write_status( + args.status, + 1.0, + f"Exported {len(result_paths)} {args.export_format} file(s)", + ) + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) + + +def parser(): + root = argparse.ArgumentParser() + commands = root.add_subparsers(dest="command", required=True) + + download = commands.add_parser("download") + download.add_argument("--repository", required=True) + download.add_argument("--model-dir", required=True) + download.add_argument("--status", required=True) + + reconstruction = commands.add_parser("reconstruct") + reconstruction.add_argument("--input", required=True) + reconstruction.add_argument("--model-dir", required=True) + reconstruction.add_argument("--output-dir", required=True) + reconstruction.add_argument( + "--export-format", + choices=EXPORT_FORMATS, + required=True, + ) + reconstruction.add_argument("--status", required=True) + reconstruction.add_argument("--process-res", type=int, required=True) + reconstruction.add_argument( + "--process-res-method", + choices=("upper_bound_resize", "lower_bound_resize"), + required=True, + ) + reconstruction.add_argument( + "--conf-thresh-percentile", + type=float, + required=True, + ) + reconstruction.add_argument("--num-max-points", type=int, required=True) + reconstruction.add_argument( + "--show-cameras", + type=int, + choices=(0, 1), + required=True, + ) + reconstruction.add_argument( + "--gs-views-interval", + type=int, + required=True, + ) + reconstruction.add_argument("--video-frames", type=int, required=True) + reconstruction.add_argument( + "--device", + choices=("auto", "cuda", "cpu"), + required=True, + ) + reconstruction.add_argument( + "--use-ray-pose", + type=int, + choices=(0, 1), + required=True, + ) + reconstruction.add_argument( + "--ref-view-strategy", + choices=("first", "middle", "saddle_balanced", "saddle_sim_range"), + required=True, + ) + return root + + +def main(): + args = parser().parse_args() + try: + if args.command == "download": + download_model(args) + else: + reconstruct(args) + except Exception as error: + write_status(args.status, 1.0, "DA3 task failed", str(error)) + raise + + +if __name__ == "__main__": + main() diff --git a/tests/test_addon.py b/tests/test_addon.py new file mode 100644 index 0000000..fc97c70 --- /dev/null +++ b/tests/test_addon.py @@ -0,0 +1,322 @@ +import importlib +import inspect +import sys +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory +from types import ModuleType, SimpleNamespace +from unittest.mock import patch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +SOURCE_ROOT = PROJECT_ROOT / "src" + + +class AddonTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.remove_modules() + fake_bpy = ModuleType("bpy") + fake_bpy.props = SimpleNamespace( + BoolProperty=lambda **options: options, + EnumProperty=lambda **options: options, + FloatProperty=lambda **options: options, + IntProperty=lambda **options: options, + PointerProperty=lambda **options: options, + StringProperty=lambda **options: options, + ) + fake_bpy.types = SimpleNamespace( + AddonPreferences=object, + Operator=object, + Panel=object, + PropertyGroup=object, + Scene=type("Scene", (), {}), + STATUSBAR_HT_header=SimpleNamespace( + prepend=lambda draw: cls.status_bar_draws.append(draw), + remove=lambda draw: cls.status_bar_draws.remove(draw), + ), + ) + cls.registered = [] + cls.status_bar_draws = [] + fake_bpy.utils = SimpleNamespace( + register_class=cls.registered.append, + unregister_class=cls.registered.remove, + ) + fake_bpy.context = SimpleNamespace( + preferences=SimpleNamespace(addons={}), + ) + fake_bpy.app = SimpleNamespace(online_access=True) + fake_bpy.path = SimpleNamespace(abspath=lambda path: path) + sys.modules["bpy"] = fake_bpy + sys.path.insert(0, str(SOURCE_ROOT)) + cls.addon = importlib.import_module("omoo_da3") + cls.fake_bpy = fake_bpy + + @classmethod + def tearDownClass(cls): + sys.path.remove(str(SOURCE_ROOT)) + sys.modules.pop("bpy", None) + cls.remove_modules() + + @staticmethod + def remove_modules(): + for module_name in tuple(sys.modules): + if module_name == "omoo_da3" or module_name.startswith("omoo_da3."): + sys.modules.pop(module_name) + + def tearDown(self): + self.registered.clear() + self.status_bar_draws.clear() + self.addon.operators._active_job = None + if hasattr(self.fake_bpy.types.Scene, "omoo_da3_settings"): + del self.fake_bpy.types.Scene.omoo_da3_settings + + def test_uses_extension_manifest_instead_of_legacy_bl_info(self): + self.assertFalse(hasattr(self.addon, "bl_info")) + + def test_registers_and_unregisters_extension_types(self): + self.addon.register() + + self.assertEqual(self.registered, list(self.addon.CLASSES)) + self.assertTrue( + hasattr(self.fake_bpy.types.Scene, "omoo_da3_settings") + ) + self.assertEqual( + self.status_bar_draws, + [self.addon.draw_status_bar], + ) + + self.addon.unregister() + + self.assertEqual(self.registered, []) + self.assertFalse( + hasattr(self.fake_bpy.types.Scene, "omoo_da3_settings") + ) + self.assertEqual(self.status_bar_draws, []) + + def test_preferences_use_package_identifier(self): + self.assertEqual( + self.addon.OMOO_DA3_preferences.bl_idname, + "omoo_da3", + ) + + def test_default_storage_is_shared_in_user_home(self): + runtime = importlib.import_module("omoo_da3.runtime") + + self.assertEqual(runtime.DEFAULT_STORAGE_ROOT.name, ".omoo_da3") + self.assertEqual( + self.addon.OMOO_DA3_preferences.__annotations__["storage_root"][ + "default" + ], + str(runtime.DEFAULT_STORAGE_ROOT), + ) + + def test_reconstruct_command_preserves_da3_parameter_names(self): + runtime = importlib.import_module("omoo_da3.runtime") + settings = SimpleNamespace( + input_path="input.png", + model="GIANT_1_1", + gs_model="NESTED_GIANT_LARGE_1_1", + process_res=504, + process_res_method="upper_bound_resize", + conf_thresh_percentile=40.0, + num_max_points=1_000_000, + show_cameras=True, + gs_views_interval=1, + video_frames=16, + device="CUDA", + use_ray_pose=False, + ref_view_strategy="saddle_balanced", + ) + with patch.object( + runtime, + "runtime_python", + return_value=Path("python"), + ), patch.object( + runtime, + "model_directory", + return_value=Path("model"), + ): + command = runtime.reconstruct_command( + settings, + "gs_ply", + Path("status.json"), + Path("result"), + ) + + self.assertIn("--export-format", command) + self.assertIn("gs_ply", command) + for parameter in ( + "--process-res", + "--process-res-method", + "--conf-thresh-percentile", + "--num-max-points", + "--show-cameras", + "--gs-views-interval", + "--use-ray-pose", + "--ref-view-strategy", + ): + with self.subTest(parameter=parameter): + self.assertIn(parameter, command) + + def test_export_format_property_includes_all_da3_formats(self): + export_format = self.addon.OMOO_DA3_PG_settings.__annotations__[ + "export_format" + ] + identifiers = [item[0] for item in export_format["items"]] + + self.assertEqual( + identifiers, + ["glb", "gs_ply", "depth_vis", "mini_npz"], + ) + self.assertEqual(export_format["default"], "glb") + + def test_gs_model_property_only_includes_gaussian_models(self): + gs_model = self.addon.OMOO_DA3_PG_settings.__annotations__[ + "gs_model" + ] + identifiers = [item[0] for item in gs_model["items"]] + + self.assertEqual( + identifiers, + ["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"], + ) + + def test_debug_panel_is_closed_by_default(self): + self.assertIn( + "DEFAULT_CLOSED", + self.addon.OMOO_DA3_PT_debug.bl_options, + ) + + def test_debug_panel_uses_one_export_button_and_dynamic_settings(self): + draw_source = inspect.getsource( + self.addon.OMOO_DA3_PT_debug.draw + ) + + self.assertEqual(draw_source.count("layout.operator("), 1) + self.assertIn('settings.export_format == "glb"', draw_source) + self.assertIn('settings.export_format == "gs_ply"', draw_source) + self.assertIn('"gs_model" if settings.export_format', draw_source) + self.assertIn("model_row.prop(settings, model_property)", draw_source) + self.assertIn('model_row.label(text="", icon="CHECKMARK")', draw_source) + self.assertIn('model_row.label(text="", icon="ERROR")', draw_source) + self.assertNotIn('text="Ready"', draw_source) + self.assertNotIn('text="Missing"', draw_source) + self.assertNotIn("License:", draw_source) + self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source) + + def test_runtime_models_panel_assigns_each_download_target(self): + draw_source = inspect.getsource( + self.addon.OMOO_DA3_PT_runtime_models.draw + ) + model_source = inspect.getsource( + self.addon.OMOO_DA3_PT_runtime_models._draw_model + ) + + self.assertIn("for model in MODELS.values()", draw_source) + self.assertIn("operator.model_key = model.key", model_source) + self.assertNotIn("Isolated Runtime ready", draw_source) + self.assertNotIn("License:", model_source) + self.assertEqual( + self.addon.OMOO_DA3_PT_runtime_models.bl_label, + "Setup", + ) + + def test_completed_export_opens_job_directory(self): + calls = [] + self.fake_bpy.ops = SimpleNamespace( + wm=SimpleNamespace( + path_open=lambda **options: calls.append(options) or {"FINISHED"}, + ), + ) + with TemporaryDirectory() as directory: + self.addon.operators.open_job_directory(directory) + + self.assertEqual(len(calls), 1) + self.assertIn("filepath", calls[0]) + + def test_status_bar_draws_progress_and_cancel_button(self): + progress_calls = [] + operator_calls = [] + + class Layout: + ui_units_x = 0.0 + + def row(self, align=False): + return self + + def label(self, **_options): + return None + + def progress(self, **options): + progress_calls.append(options) + + def operator(self, operator, **options): + operator_calls.append((operator, options)) + + self.addon.operators._active_job = SimpleNamespace() + context = SimpleNamespace( + scene=SimpleNamespace( + omoo_da3_settings=SimpleNamespace( + status="Processing", + progress=0.25, + ), + ), + ) + + self.addon.draw_status_bar( + SimpleNamespace(layout=Layout()), + context, + ) + + self.assertEqual( + progress_calls, + [{"factor": 0.25, "type": "BAR", "text": "25%"}], + ) + self.assertEqual( + operator_calls, + [("omoo_da3.cancel_job", {"text": "", "icon": "X"})], + ) + + def test_ui_redraw_targets_status_bar_and_view_3d_areas(self): + redrawn = [] + areas = ( + SimpleNamespace( + type="STATUSBAR", + tag_redraw=lambda: redrawn.append("status"), + ), + SimpleNamespace( + type="VIEW_3D", + tag_redraw=lambda: redrawn.append("view"), + ), + ) + context = SimpleNamespace( + window_manager=SimpleNamespace( + windows=( + SimpleNamespace(screen=SimpleNamespace(areas=areas)), + ), + ), + ) + + self.addon.status_bar.redraw_ui(context) + + self.assertEqual(redrawn, ["status", "view"]) + + def test_forced_ui_redraw_swaps_window_immediately(self): + calls = [] + self.fake_bpy.ops = SimpleNamespace( + wm=SimpleNamespace( + redraw_timer=lambda **options: calls.append(options), + ), + ) + context = SimpleNamespace( + window=SimpleNamespace(), + window_manager=SimpleNamespace(windows=()), + ) + + self.addon.status_bar.redraw_ui(context, force=True) + + self.assertEqual( + calls, + [{"type": "DRAW_WIN_SWAP", "iterations": 1}], + ) diff --git a/tests/test_bootstrap_runtime.py b/tests/test_bootstrap_runtime.py new file mode 100644 index 0000000..8991b7f --- /dev/null +++ b/tests/test_bootstrap_runtime.py @@ -0,0 +1,63 @@ +import importlib.util +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def load_bootstrap_runtime(): + specification = importlib.util.spec_from_file_location( + "omoo_da3_bootstrap_runtime", + PROJECT_ROOT / "src" / "omoo_da3" / "bootstrap_runtime.py", + ) + module = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + +class BootstrapRuntimeTest(unittest.TestCase): + def test_installs_da3_and_non_compiled_gsplat(self): + bootstrap = load_bootstrap_runtime() + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + runtime_dir = root / "runtime" + runtime_dir.mkdir() + args = SimpleNamespace( + uv="uv", + runtime_dir=str(runtime_dir), + backend="cu128", + da3_commit="test-commit", + status=str(root / "status.json"), + ) + with patch.object(bootstrap, "run") as mocked_run: + bootstrap.install(args) + + commands = [call.args[0] for call in mocked_run.call_args_list] + da3_command = next( + command + for command in commands + if any("depth-anything-3 @" in str(value) for value in command) + ) + self.assertIn( + f"addict=={bootstrap.ADDICT_VERSION}", + da3_command, + ) + gsplat_call = next( + call + for call in mocked_run.call_args_list + if any("gsplat @" in str(value) for value in call.args[0]) + ) + gsplat_command = gsplat_call.args[0] + self.assertIn("--no-build-isolation", gsplat_command) + self.assertTrue( + any(bootstrap.GSPLAT_COMMIT in str(value) for value in gsplat_command) + ) + self.assertEqual(gsplat_call.args[1], {"BUILD_NO_CUDA": "1"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_manifest.py b/tests/test_manifest.py new file mode 100644 index 0000000..7fcf356 --- /dev/null +++ b/tests/test_manifest.py @@ -0,0 +1,25 @@ +import tomllib +import unittest +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +MANIFEST_PATH = PROJECT_ROOT / "src" / "omoo_da3" / "blender_manifest.toml" + + +class ManifestTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.manifest = tomllib.loads(MANIFEST_PATH.read_text(encoding="utf-8")) + + def test_declares_modern_extension_metadata(self): + self.assertEqual(self.manifest["schema_version"], "1.0.0") + self.assertEqual(self.manifest["id"], "omoo_da3") + self.assertEqual(self.manifest["type"], "add-on") + self.assertNotIn("bl_info", self.manifest) + + def test_declares_required_runtime_permissions(self): + self.assertEqual( + set(self.manifest["permissions"]), + {"network", "files"}, + ) diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..de693b1 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,45 @@ +import importlib.util +import unittest +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def load_module(name, relative_path): + specification = importlib.util.spec_from_file_location( + name, + PROJECT_ROOT / relative_path, + ) + module = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + +class ModelTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.models = load_module( + "omoo_da3_models", + "src/omoo_da3/models.py", + ) + + def test_base_model_is_default_friendly_apache_model(self): + model = self.models.get_model("BASE") + + self.assertEqual(model.repository, "depth-anything/DA3-BASE") + self.assertEqual(model.license_name, "Apache 2.0") + + def test_large_model_exposes_noncommercial_license(self): + model = self.models.get_model("LARGE_1_1") + + self.assertEqual(model.license_name, "CC BY-NC 4.0") + + def test_giant_models_expose_gaussian_head(self): + for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"): + with self.subTest(model_key=model_key): + self.assertTrue(self.models.get_model(model_key).supports_gs) + + def test_unknown_model_is_rejected(self): + with self.assertRaisesRegex(ValueError, "Unknown DA3 model"): + self.models.get_model("UNKNOWN") diff --git a/tests/test_pack.py b/tests/test_pack.py new file mode 100644 index 0000000..36cbd6e --- /dev/null +++ b/tests/test_pack.py @@ -0,0 +1,39 @@ +import importlib.util +import tempfile +import unittest +from pathlib import Path +from zipfile import ZipFile + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def load_pack(): + specification = importlib.util.spec_from_file_location( + "omoo_da3_pack", + PROJECT_ROOT / "pack.py", + ) + module = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + +class PackTest(unittest.TestCase): + def test_builds_extension_root_layout(self): + pack = load_pack() + with tempfile.TemporaryDirectory() as directory: + original_dist = pack.DIST_DIR + pack.DIST_DIR = Path(directory) + try: + output_path = pack.build_extension() + finally: + pack.DIST_DIR = original_dist + + with ZipFile(output_path) as archive: + names = set(archive.namelist()) + + self.assertIn("__init__.py", names) + self.assertIn("blender_manifest.toml", names) + self.assertIn("worker.py", names) + self.assertNotIn("importer.py", names) + self.assertNotIn("omoo_da3/__init__.py", names) diff --git a/tests/test_worker.py b/tests/test_worker.py new file mode 100644 index 0000000..95267d3 --- /dev/null +++ b/tests/test_worker.py @@ -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() diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..9825743 --- /dev/null +++ b/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 3 +requires-python = ">=3.11" + +[[package]] +name = "omoo-da3" +version = "0.8.1" +source = { virtual = "." }