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

feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe

Browse Source 
- Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing.
- Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations.
- Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Added model synchronization script `sync_models.py` for downloading and verifying ONNX models.
- Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality.
- Updated `vendor.lock.toml` to remove unused artifacts.
This commit is contained in:
马南
2026-08-14 10:27:09 +08:00
parent 13fab9ce88
commit 2ad30bf4a3
41 changed files with 1569 additions and 2109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/snapform/__init__.py
View File

@@ -21,9 +21,9 @@ from .operators import (
)
from .panel import (
Ben2RemovalDebugPanel,
AiServicePanel,
Da3OutputDebugPanel,
DebugPanel,
AiServicePanel,
ModelsPanel,
Moge2ApiDebugPanel,
)
@@ -50,8 +50,8 @@ CLASSES = (
AiServicePanel,
ModelsPanel,
DebugPanel,
Da3OutputDebugPanel,
Moge2ApiDebugPanel,
Da3OutputDebugPanel,
Ben2RemovalDebugPanel,
)

50
src/snapform/ben2_onnx.py Normal file
View File

@@ -0,0 +1,50 @@
"""BEN2 ONNX inference with Pillow and NumPy preprocessing."""
from pathlib import Path
try:
from .onnx_runtime import create_session as create_runtime_session
except ImportError:
from onnx_runtime import create_session as create_runtime_session
ONNX_PATH = Path("onnx/model_fp16.onnx")
INPUT_SIZE = (1024, 1024)
def model_path(model_directory):
return Path(model_directory) / ONNX_PATH
def create_session(model_directory, requested_device="auto"):
return create_runtime_session(model_path(model_directory), requested_device)
def preprocess(image):
import numpy as np
from PIL import Image
resized = image.convert("RGB").resize(INPUT_SIZE, Image.Resampling.LANCZOS)
array = np.asarray(resized, dtype=np.float32) / 255.0
mean = np.array((0.485, 0.456, 0.406), dtype=np.float32)
std = np.array((0.229, 0.224, 0.225), dtype=np.float32)
array = (array - mean) / std
return np.ascontiguousarray(array.transpose(2, 0, 1)[None])
def infer(session, image):
import numpy as np
from PIL import Image
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
alpha = session.run([output_name], {input_name: preprocess(image)})[0]
alpha = np.asarray(alpha, dtype=np.float32).squeeze()
minimum = float(alpha.min())
span = max(float(alpha.max()) - minimum, np.finfo(np.float32).eps)
alpha = np.clip((alpha - minimum) / span * 255.0, 0.0, 255.0)
mask = Image.fromarray(alpha.astype(np.uint8), mode="L")
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
foreground = image.convert("RGBA")
foreground.putalpha(mask)
return foreground

16
src/snapform/client.py
View File

@@ -26,10 +26,8 @@ from .models import (
)
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
MANIFEST_FILENAME = "manifest.json"
MANIFEST_REVISION = 5
MANIFEST_REVISION = 7
SERVICE_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = {
@@ -172,8 +170,8 @@ def env_devices():
def machine_devices():
"""Devices the machine can run before environment installation."""
devices = ["cpu"]
if shutil.which("nvidia-smi"):
devices.insert(0, "cuda")
if sys.platform == "win32":
devices.insert(0, "directml")
elif sys.platform == "darwin" and platform.machine().lower() in {
"arm64",
"aarch64",
@@ -189,6 +187,7 @@ def supported_devices():
DEVICE_LABELS = {
"cuda": ("CUDA", "NVIDIA CUDA"),
"directml": ("DirectML", "Windows GPU (DirectML)"),
"mps": ("MPS", "Apple Metal (MPS)"),
"cpu": ("CPU", "CPU"),
}
@@ -283,7 +282,6 @@ def uv_cache_path():
def install_env_command(status_path):
executable = uv_executable()
torch_variant = "cu128" if "cuda" in machine_devices() else "cpu"
script = Path(__file__).with_name("install_env.py")
command = [
sys.executable,
@@ -296,12 +294,6 @@ def install_env_command(status_path):
str(storage_root() / MANIFEST_FILENAME),
"--models-dir",
str(models_directory()),
"--torch-variant",
torch_variant,
"--da3-commit",
DA3_COMMIT,
"--ben2-commit",
BEN2_COMMIT,
"--status",
str(status_path),
]

221
src/snapform/da3_onnx.py Normal file
View File

@@ -0,0 +1,221 @@
"""Depth Anything 3 inference with ONNX Runtime and NumPy."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import numpy as np
try:
from .onnx_runtime import create_session as create_onnx_session
except ImportError:
from onnx_runtime import create_session as create_onnx_session
ONNX_PATH = "model.onnx"
PATCH_SIZE = 14
STRATEGY_CODES = {
"first": 0,
"middle": 1,
"saddle_balanced": 2,
"saddle_sim_range": 3,
}
IMAGE_MEAN = np.asarray((0.485, 0.456, 0.406), dtype=np.float32)
IMAGE_STD = np.asarray((0.229, 0.224, 0.225), dtype=np.float32)
@dataclass
class Prediction:
depth: np.ndarray
conf: np.ndarray
extrinsics: np.ndarray
intrinsics: np.ndarray
processed_images: np.ndarray
is_metric: int = 0
sky: np.ndarray | None = None
gaussians: object | None = None
aux: dict | None = None
scale_factor: float | None = None
def create_session(model_dir, requested_device="auto"):
path = Path(model_dir) / ONNX_PATH
return create_onnx_session(path, requested_device)
def nearest_multiple(value, multiple=PATCH_SIZE):
lower = value // multiple * multiple
upper = lower + multiple
return upper if abs(upper - value) <= abs(value - lower) else lower
def target_size(size, process_res, method):
width, height = size
boundary = max(width, height)
if method.startswith("lower_bound"):
boundary = min(width, height)
scale = float(process_res) / boundary
width = max(1, int(round(width * scale)))
height = max(1, int(round(height * scale)))
if method.endswith("resize"):
return (
max(PATCH_SIZE, nearest_multiple(width)),
max(PATCH_SIZE, nearest_multiple(height)),
)
if method.endswith("crop"):
return (
max(PATCH_SIZE, width // PATCH_SIZE * PATCH_SIZE),
max(PATCH_SIZE, height // PATCH_SIZE * PATCH_SIZE),
)
raise ValueError(f"Unsupported process_res_method: {method}")
def resize_and_crop(image, process_res, method):
from PIL import Image
width, height = image.size
boundary = max(width, height)
if method.startswith("lower_bound"):
boundary = min(width, height)
scale = float(process_res) / boundary
resized_size = (
max(1, int(round(width * scale))),
max(1, int(round(height * scale))),
)
resampling = (
Image.Resampling.BICUBIC
if scale > 1.0
else Image.Resampling.BOX
)
image = image.resize(resized_size, resample=resampling)
final_size = target_size((width, height), process_res, method)
if method.endswith("resize"):
if image.size != final_size:
upscale = final_size[0] > image.width or final_size[1] > image.height
resampling = (
Image.Resampling.BICUBIC
if upscale
else Image.Resampling.BOX
)
image = image.resize(final_size, resample=resampling)
return image
left = (image.width - final_size[0]) // 2
top = (image.height - final_size[1]) // 2
return image.crop((left, top, left + final_size[0], top + final_size[1]))
def preprocess(image_paths, process_res=504, process_res_method="upper_bound_resize"):
from PIL import Image
processed = []
for path in image_paths:
with Image.open(path) as source:
image = resize_and_crop(
source.convert("RGB"),
process_res,
process_res_method,
)
processed.append(np.asarray(image, dtype=np.uint8))
min_height = min(image.shape[0] for image in processed)
min_width = min(image.shape[1] for image in processed)
cropped = []
for image in processed:
top = (image.shape[0] - min_height) // 2
left = (image.shape[1] - min_width) // 2
cropped.append(
image[top : top + min_height, left : left + min_width]
)
images = np.stack(cropped)
normalized = images.astype(np.float32) / 255.0
normalized = (normalized - IMAGE_MEAN) / IMAGE_STD
model_input = normalized.transpose(0, 3, 1, 2)[None]
return np.ascontiguousarray(model_input), images
def quaternion_to_matrix(quaternion):
i, j, k, real = np.moveaxis(quaternion, -1, 0)
two_scale = 2.0 / np.sum(quaternion * quaternion, axis=-1)
values = (
1 - two_scale * (j * j + k * k),
two_scale * (i * j - k * real),
two_scale * (i * k + j * real),
two_scale * (i * j + k * real),
1 - two_scale * (i * i + k * k),
two_scale * (j * k - i * real),
two_scale * (i * k - j * real),
two_scale * (j * k + i * real),
1 - two_scale * (i * i + j * j),
)
return np.stack(values, axis=-1).reshape(quaternion.shape[:-1] + (3, 3))
def camera_from_pose(pose_encoding, height, width):
translation = pose_encoding[..., :3]
rotation = quaternion_to_matrix(pose_encoding[..., 3:7])
rotation_inverse = np.swapaxes(rotation, -1, -2)
translation_inverse = -rotation_inverse @ translation[..., None]
extrinsics = np.concatenate(
(rotation_inverse, translation_inverse),
axis=-1,
).astype(np.float32)
field_of_view_height = pose_encoding[..., 7]
field_of_view_width = pose_encoding[..., 8]
focal_y = (height / 2.0) / np.maximum(
np.tan(field_of_view_height / 2.0),
1e-6,
)
focal_x = (width / 2.0) / np.maximum(
np.tan(field_of_view_width / 2.0),
1e-6,
)
intrinsics = np.zeros(pose_encoding.shape[:-1] + (3, 3), dtype=np.float32)
intrinsics[..., 0, 0] = focal_x
intrinsics[..., 1, 1] = focal_y
intrinsics[..., 0, 2] = width / 2.0
intrinsics[..., 1, 2] = height / 2.0
intrinsics[..., 2, 2] = 1.0
return extrinsics, intrinsics
def infer(
session,
image_paths,
process_res=504,
process_res_method="upper_bound_resize",
ref_view_strategy="saddle_balanced",
use_ray_pose=False,
):
del use_ray_pose
inputs = session.get_inputs()
model_input, processed_images = preprocess(
image_paths,
process_res,
process_res_method,
)
strategy_code = np.asarray(
[STRATEGY_CODES.get(ref_view_strategy, 2)],
dtype=np.int64,
)
input_names = [value.name for value in inputs]
outputs = session.run(
None,
{
input_names[0]: model_input,
input_names[1]: strategy_code,
},
)
depth, confidence, pose_encoding = outputs[:3]
height, width = depth.shape[-2:]
extrinsics, intrinsics = camera_from_pose(
pose_encoding[0],
height,
width,
)
return Prediction(
depth=np.asarray(depth[0], dtype=np.float32),
conf=np.asarray(confidence[0], dtype=np.float32),
extrinsics=extrinsics,
intrinsics=intrinsics,
processed_images=processed_images,
aux={},
)

12
src/snapform/image_depth.py
View File

@@ -2239,12 +2239,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
name="Depth From",
description="Model or original image used to build depth",
items=image_depth_model_items,
default=1,
default=0,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Use ray-based camera pose estimation",
default=True,
description="Legacy task compatibility; ONNX uses the camera decoder",
default=False,
)
snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_resolution_level: bpy.props.IntProperty(
@@ -2278,14 +2278,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
layout = self.layout
layout.prop(self, "snapform_model")
if self.snapform_model in {"SMALL", "BASE"}:
layout.prop(self, "snapform_use_ray_pose")
draw_process_resolution(layout, self, context.object.data)
elif self.snapform_model.startswith("MOGE2_"):
from .client import model_ready
if not model_ready(self.snapform_model):
layout.label(text="MoGe-2 model is not installed", icon="ERROR")
layout.label(text="Download it in Models")
layout.prop(self, "snapform_moge_resolution_level")
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_unshaded")

7
src/snapform/image_mesh.py
View File

@@ -608,12 +608,12 @@ class ConvertImageToMesh(bpy.types.Operator):
name="Depth From",
description="Model used to estimate image depth",
items=image_surface_model_items,
default=1,
default=0,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Use ray-based camera pose estimation",
default=True,
description="Legacy task compatibility; ONNX uses the camera decoder",
default=False,
)
snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_resolution_level: bpy.props.IntProperty(
@@ -670,7 +670,6 @@ class ConvertImageToMesh(bpy.types.Operator):
if self.snapform_model.startswith("MOGE2_"):
layout.prop(self, "snapform_moge_resolution_level")
else:
layout.prop(self, "snapform_use_ray_pose")
draw_process_resolution(layout, self, context.object.data)
layout.prop(self, "snapform_create_camera")
layout.prop(self, "snapform_remove_background")

262
src/snapform/install_env.py
View File

@@ -1,5 +1,4 @@
import argparse
import filecmp
import json
import os
import platform
@@ -28,47 +27,36 @@ PYPI_MIRRORS = (
"https://mirrors.aliyun.com/pypi/simple/",
)
PYPI_DEFAULT_INDEX = "https://pypi.org/simple/"
TORCH_VERSION = "2.9.1"
TORCHVISION_VERSION = "0.24.1"
XFORMERS_VERSION = "0.0.33.post2"
ADDICT_VERSION = "2.4.0"
FASTAPI_VERSION = "0.139.2"
UVICORN_VERSION = "0.51.0"
MANIFEST_REVISION = 5
TORCH_INDEXES = {
"cu128": "https://download.pytorch.org/whl/cu128",
"cpu": "https://download.pytorch.org/whl/cpu",
}
ONNXRUNTIME_VERSION = "1.24.4"
NUMPY_VERSION = "2.4.2"
PILLOW_VERSION = "12.1.1"
OPENCV_VERSION = "4.13.0.92"
TRIMESH_VERSION = "4.11.2"
HUGGINGFACE_HUB_VERSION = "1.4.1"
MANIFEST_REVISION = 7
_resolved_pypi_index = None
DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key)
RUNTIME_STAGE_COUNT = 3
RUNTIME_STAGE_COUNT = 2
ENVIRONMENT_SLICE = (0.02, 0.06)
PYTORCH_SLICE = (0.06, 0.48)
XFORMERS_SLICE = (0.48, 0.62)
AI_COMPONENTS_SLICE = (0.62, 0.85)
AI_COMPONENTS_SLICE = (0.10, 0.85)
MODELS_SLICE_START = 0.85
MODELS_SLICE_END = 1.0
RAMP_RATE = 0.02
CREEP_CAP = 0.60
INSTALL_ESTIMATES = {
"PyTorch": 8 * 60,
"AI components": 8 * 60,
"ONNX Runtime": 3 * 60,
}
INSTALL_ESTIMATE_RANGES = {
"PyTorch": PYTORCH_SLICE,
"AI components": (XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1]),
"ONNX Runtime": AI_COMPONENTS_SLICE,
}
CHECK_CODE = (
"import json; import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"from moge.model.v2 import MoGeModel; "
"import fastapi, uvicorn; "
"mps = getattr(torch.backends, 'mps', None); "
"import json; import fastapi, uvicorn, onnxruntime; "
"import cv2, numpy, PIL, trimesh; "
"print('SNAPFORM_DEVICES=' + json.dumps({"
"'torch': torch.__version__, "
"'cuda': torch.cuda.is_available(), "
"'mps': bool(mps is not None and mps.is_available())}))"
"'onnxruntime': onnxruntime.__version__, "
"'providers': onnxruntime.get_available_providers()}))"
)
@@ -158,10 +146,6 @@ def estimated_remaining(label, progress, elapsed=0.0):
return max(expected_seconds * remaining_fraction - elapsed, 0.0)
def effective_torch_index(torch_variant):
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(torch_variant)
def resolve_indexes(args):
"""???????????????????????????????????"""
global _resolved_pypi_index
@@ -320,52 +304,13 @@ def venv_python(venv_dir):
return Path(venv_dir) / "bin" / "python"
def install_stages(args, venv_dir, python, uv, vendor_dir=None):
# macOS 的 PyTorch wheel 只在 PyPI 上，download.pytorch.org 没有 macosx 构建
torch_index = (
None
if sys.platform == "darwin"
else effective_torch_index(args.torch_variant)
def install_stages(args, venv_dir, python, uv):
"""Build the lightweight ONNX-only runtime and model download stages."""
runtime_package = (
f"onnxruntime-directml=={ONNXRUNTIME_VERSION}"
if sys.platform == "win32"
else f"onnxruntime=={ONNXRUNTIME_VERSION}"
)
vendor_dir = Path(vendor_dir or Path(venv_dir).parent / "vendor")
def archive_requirement(project_name, filename):
local = Path(vendor_dir) / filename
if not local.is_file():
raise RuntimeError(
"Bundled vendor artifact is missing; reinstall the extension: "
f"{local}"
)
return f"{project_name} @ {local.as_uri()}"
da3_requirement = archive_requirement(
"depth-anything-3",
"depth-anything-3.tar.gz",
)
ben2_requirement = archive_requirement(
"ben2",
"ben2.tar.gz",
)
utils3d_requirement = archive_requirement(
"utils3d",
"utils3d.tar.gz",
)
moge_requirement = archive_requirement(
"moge",
"moge.tar.gz",
)
torch_command = [
uv,
"pip",
"install",
"-v",
"--python",
python,
f"torch=={TORCH_VERSION}",
f"torchvision=={TORCHVISION_VERSION}",
]
if torch_index is not None:
torch_command.extend(["--index-url", torch_index])
stages = [
(
"Environment",
@@ -374,13 +319,7 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None):
"plain",
),
(
"PyTorch",
torch_command,
None,
"install",
),
(
"AI components",
"ONNX Runtime",
[
uv,
"pip",
@@ -388,38 +327,19 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None):
"-v",
"--python",
python,
da3_requirement,
f"addict=={ADDICT_VERSION}",
ben2_requirement,
utils3d_requirement,
moge_requirement,
f"fastapi=={FASTAPI_VERSION}",
f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}",
f"numpy=={NUMPY_VERSION}",
f"opencv-python-headless=={OPENCV_VERSION}",
runtime_package,
f"pillow=={PILLOW_VERSION}",
f"trimesh=={TRIMESH_VERSION}",
f"uvicorn=={UVICORN_VERSION}",
"setuptools",
"wheel",
],
None,
"install",
),
]
if sys.platform != "darwin":
xformers_command = [
uv,
"pip",
"install",
"-v",
"--python",
python,
f"xformers=={XFORMERS_VERSION}",
]
if args.torch_variant == "cu128":
xformers_command.extend(["--index-url", torch_index])
else:
xformers_command.append("--no-deps")
stages.insert(
2,
("AI components", xformers_command, None, "install"),
)
models = DOWNLOADABLE_MODELS
worker_script = Path(__file__).with_name("worker.py")
models_dir = Path(args.models_dir)
@@ -475,13 +395,13 @@ def env_installed(venv_dir, manifest_path=None):
)
def installed_devices(args):
def installed_devices():
devices = []
machine = platform.machine().lower()
if sys.platform == "darwin" and machine in {"arm64", "aarch64"}:
devices.append("mps")
elif args.torch_variant.startswith("cu"):
devices.append("cuda")
elif sys.platform == "win32":
devices.append("directml")
devices.append("cpu")
return devices
@@ -592,82 +512,6 @@ def bundled_uv_archive():
return path if path.is_file() else None
def vendor_overlay_directory(filename):
names = {
"depth-anything-3.tar.gz": "depth-anything-3",
"moge.tar.gz": "moge",
"utils3d.tar.gz": "utils3d",
}
overlay_name = names.get(filename)
if overlay_name is None:
return None
project_root = Path(__file__).resolve().parents[2]
overlay = project_root / "patches" / overlay_name
return overlay if overlay.is_dir() else None
def copy_archive_with_overlay(source, target, overlay):
"""开发目录使用原始 vendor 时，应用与正式打包相同的源码覆盖。"""
with tarfile.open(source, "r:gz") as input_archive:
members = input_archive.getmembers()
top_level = members[0].name.split("/", 1)[0]
with tarfile.open(target, "w:gz") as output_archive:
for member in members:
if member.isdir() or member.islnk() or member.issym():
output_archive.addfile(member)
continue
relative = member.name[len(top_level) + 1 :]
replacement = overlay / relative
if replacement.is_file():
info = tarfile.TarInfo(member.name)
info.size = replacement.stat().st_size
info.mode = member.mode
with replacement.open("rb") as stream:
output_archive.addfile(info, stream)
continue
with input_archive.extractfile(member) as stream:
output_archive.addfile(member, stream)
def copy_vendor_inputs(bundled_dir, vendor_dir):
"""Copy bundled installation inputs into a temporary directory."""
for name in (
"depth-anything-3.tar.gz",
"ben2.tar.gz",
"utils3d.tar.gz",
"moge.tar.gz",
):
bundled = Path(bundled_dir) / name
if not bundled.is_file():
continue
target = Path(vendor_dir) / name
overlay = vendor_overlay_directory(name)
if overlay is None and target.is_file() and filecmp.cmp(
bundled,
target,
shallow=False,
):
continue
Path(vendor_dir).mkdir(parents=True, exist_ok=True)
if overlay is not None and tarfile.is_tarfile(bundled):
copy_archive_with_overlay(bundled, target, overlay)
else:
shutil.copy2(bundled, target)
def remove_legacy_vendor(storage_root):
"""Remove installation inputs persisted by older development builds."""
root = Path(storage_root)
shutil.rmtree(root / "vendor", ignore_errors=True)
for name in (
"depth-anything-3.tar.gz",
"ben2.tar.gz",
"utils3d.tar.gz",
"moge.tar.gz",
):
(root / name).unlink(missing_ok=True)
def archive_is_valid(path):
if zipfile.is_zipfile(path):
return True
@@ -800,14 +644,6 @@ def run_stage(
def install(args):
if sys.platform == "darwin" and platform.machine().lower() not in {
"arm64",
"aarch64",
}:
raise RuntimeError(
"SnapForm requires Apple Silicon macOS; PyTorch stopped "
"shipping Intel Mac (x86_64) wheels after 2.2.2"
)
resolve_indexes(args)
venv_dir = Path(args.venv_dir)
manifest_path = Path(
@@ -823,13 +659,7 @@ def install(args):
uv = Path(resolved) if resolved else uv
uv_missing = not uv.is_file()
python = venv_python(venv_dir)
vendor_temp = tempfile.TemporaryDirectory(prefix="snapform_vendor_")
vendor_dir = Path(vendor_temp.name)
copy_vendor_inputs(
bundled_vendor_directory(),
vendor_dir,
)
stages = install_stages(args, venv_dir, python, uv, vendor_dir)
stages = install_stages(args, venv_dir, python, uv)
env_stages = [stage for stage in stages if stage[3] != "download"]
model_stages = [stage for stage in stages if stage[3] == "download"]
model_span = (
@@ -838,23 +668,15 @@ def install(args):
else 0.0
)
installed_packages = []
component_stages = env_stages[2:]
component_ranges = (
[XFORMERS_SLICE, AI_COMPONENTS_SLICE]
if len(component_stages) == 2
else [(XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1])]
)
for index, (label, command, environment, kind) in enumerate(env_stages):
if index == 0:
stage = 1
stage_start, stage_end = ENVIRONMENT_SLICE
elif index == 1:
stage = 2
stage_start, stage_end = PYTORCH_SLICE
stage_start, stage_end = AI_COMPONENTS_SLICE
else:
stage = 3
component_index = index - 2
stage_start, stage_end = component_ranges[component_index]
raise RuntimeError(f"Unexpected runtime stage: {label}")
if uv_missing and index == 0:
midpoint = stage_start + (stage_end - stage_start) * 0.5
install_uv(
@@ -880,9 +702,7 @@ def install(args):
)
if isinstance(stage_installed, list):
installed_packages.extend(stage_installed)
vendor_temp.cleanup()
remove_legacy_vendor(venv_dir.parent)
devices = installed_devices(args)
devices = installed_devices()
venv_dir.mkdir(parents=True, exist_ok=True)
manifest_path.parent.mkdir(parents=True, exist_ok=True)
ready_path = manifest_path
@@ -891,14 +711,13 @@ def install(args):
{
"manifest_revision": MANIFEST_REVISION,
"python": "3.12",
"torch": TORCH_VERSION,
"xformers": XFORMERS_VERSION,
"onnxruntime": ONNXRUNTIME_VERSION,
"fastapi": FASTAPI_VERSION,
"uvicorn": UVICORN_VERSION,
"torch_variant": args.torch_variant,
"runtime_variant": (
"directml" if sys.platform == "win32" else "default"
),
"devices": devices,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
"packages": sorted(set(installed_packages)),
},
indent=2,
@@ -952,9 +771,6 @@ def parse_args():
parser.add_argument("--venv-dir", required=True)
parser.add_argument("--manifest", required=True)
parser.add_argument("--models-dir", required=True)
parser.add_argument("--torch-variant", choices=tuple(TORCH_INDEXES), required=True)
parser.add_argument("--da3-commit", required=True)
parser.add_argument("--ben2-commit", required=True)
parser.add_argument("--status", required=True)
parser.add_argument("--pypi-index", default=None)
return parser.parse_args()

131
src/snapform/models.py
View File

@@ -21,42 +21,33 @@ MODELS = {
"SMALL": ModelSpec(
key="SMALL",
label="DA3 Small",
directory_name="da3_small",
directory_name="DA3-SMALL",
repository="depth-anything/DA3-SMALL",
license_name="Apache 2.0",
description="Fastest any-view model with Apache 2.0 weights",
description="Fastest any-view model using ONNX Runtime",
supports_gs=False,
r2_directory="DA3Small",
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="DA3-SMALL",
r2_files=(
("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"),
("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"),
("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"),
),
),
"BASE": ModelSpec(
key="BASE",
label="DA3 Base",
directory_name="da3_base",
directory_name="DA3-BASE",
repository="depth-anything/DA3-BASE",
license_name="Apache 2.0",
description="Balanced any-view model with Apache 2.0 weights",
description="Balanced any-view model using ONNX Runtime",
supports_gs=False,
r2_directory="DA3Base",
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="DA3-BASE",
r2_files=(
("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"),
("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"),
("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"),
("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"),
),
),
"LARGE_1_1": ModelSpec(
key="LARGE_1_1",
label="DA3 Large 1.1",
directory_name="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,
enabled=False,
),
"GIANT_1_1": ModelSpec(
key="GIANT_1_1",
label="DA3 Giant 1.1",
@@ -82,52 +73,76 @@ MODELS = {
BEN2_MODEL = ModelSpec(
key="BEN2_BASE",
label="BEN2 Base",
directory_name="ben2_base",
repository="PramaLLC/BEN2",
license_name="MIT",
description="Fast high-quality background removal model",
directory_name="BEN2-ONNX",
repository="onnx-community/BEN2-ONNX",
license_name="Apache 2.0",
description="Fast high-quality FP16 ONNX background removal model",
supports_gs=False,
download_patterns=("config.json", "model.safetensors"),
r2_directory="BEN2Base",
r2_files=(
("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"),
("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"),
),
download_patterns=("onnx/model_fp16.onnx",),
ready_patterns=("onnx/model_fp16.onnx",),
r2_directory="BEN2-ONNX",
r2_files=((
"onnx/model_fp16.onnx",
219121675,
"dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1",
),),
)
def moge2_model(key, label, directory_name, repository, description, r2_directory, size, checksum):
return ModelSpec(
key=key,
label=label,
directory_name=directory_name,
repository=repository,
license_name="MIT",
description=description,
supports_gs=False,
download_patterns=("model.pt",),
ready_patterns=("model.pt",),
r2_directory=r2_directory,
r2_files=(("model.pt", size, checksum),),
)
MOGE2_MODELS = {
"MOGE2_VITS_NORMAL": moge2_model(
"MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal",
"Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal",
140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc",
"MOGE2_VITS_NORMAL": ModelSpec(
key="MOGE2_VITS_NORMAL",
label="MoGe-2 ViT-S Normal",
directory_name="moge-2-vits-normal-onnx",
repository="Ruicheng/moge-2-vits-normal-onnx",
license_name="MIT",
description="Fastest MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vits-normal-onnx",
r2_files=((
"model.onnx",
140852051,
"24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0",
),),
),
"MOGE2_VITB_NORMAL": moge2_model(
"MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal",
"Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal",
419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6",
"MOGE2_VITB_NORMAL": ModelSpec(
key="MOGE2_VITB_NORMAL",
label="MoGe-2 ViT-B Normal",
directory_name="moge-2-vitb-normal-onnx",
repository="Ruicheng/moge-2-vitb-normal-onnx",
license_name="MIT",
description="Balanced MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vitb-normal-onnx",
r2_files=((
"model.onnx",
419411850,
"bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21",
),),
),
"MOGE2_VITL_NORMAL": moge2_model(
"MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal",
"Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal",
1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01",
"MOGE2_VITL_NORMAL": ModelSpec(
key="MOGE2_VITL_NORMAL",
label="MoGe-2 ViT-L Normal",
directory_name="moge-2-vitl-normal-onnx",
repository="Ruicheng/moge-2-vitl-normal-onnx",
license_name="MIT",
description="Highest-quality MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vitl-normal-onnx",
r2_files=((
"model.onnx",
1324265014,
"afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
),),
),
}
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITB_NORMAL"]
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"]
DOWNLOADABLE_MODELS = {
**MODELS,
BEN2_MODEL.key: BEN2_MODEL,

165
src/snapform/moge_onnx.py Normal file
View File

@@ -0,0 +1,165 @@
"""MoGe-2 ONNX inference and NumPy post-processing."""
from pathlib import Path
import sys
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
try:
from .onnx_runtime import create_session as create_runtime_session
from .onnx_runtime import provider_label, select_providers
except ImportError:
from onnx_runtime import create_session as create_runtime_session
from onnx_runtime import provider_label, select_providers
ONNX_FILENAME = "model.onnx"
MIN_TOKENS = 1200
MAX_TOKENS = 3600
def is_onnx_model(model_directory):
return (Path(model_directory) / ONNX_FILENAME).is_file()
def create_session(model_directory, requested_device="auto"):
path = Path(model_directory) / ONNX_FILENAME
return create_runtime_session(path, requested_device)
def normalized_view_plane_uv(width, height, dtype):
import numpy as np
aspect_ratio = width / height
span_x = aspect_ratio / (1.0 + aspect_ratio**2) ** 0.5
span_y = 1.0 / (1.0 + aspect_ratio**2) ** 0.5
u = np.linspace(
-span_x * (width - 1) / width,
span_x * (width - 1) / width,
width,
dtype=dtype,
)
v = np.linspace(
-span_y * (height - 1) / height,
span_y * (height - 1) / height,
height,
dtype=dtype,
)
u, v = np.meshgrid(u, v, indexing="xy")
return np.stack((u, v), axis=-1)
def solve_focal_shift(uv, points):
import numpy as np
uv = uv.reshape(-1, 2)
xy = points[..., :2].reshape(-1, 2)
z = points[..., 2].reshape(-1)
matrix = np.stack((xy, -uv), axis=-1).reshape(-1, 2)
target = (uv * z[:, None]).reshape(-1)
focal, shift = np.linalg.lstsq(matrix, target, rcond=None)[0]
epsilon = np.finfo(np.float32).eps
for _iteration in range(20):
denominator = z + shift
denominator = np.where(
np.abs(denominator) < epsilon,
np.copysign(epsilon, denominator),
denominator,
)
projected = xy / denominator[:, None]
residual = (focal * projected - uv).reshape(-1)
jacobian = np.stack(
(
projected,
-focal * xy / np.square(denominator)[:, None],
),
axis=-1,
).reshape(-1, 2)
delta = np.linalg.lstsq(jacobian, -residual, rcond=None)[0]
focal += delta[0]
shift += delta[1]
if np.max(np.abs(delta)) < 1e-6:
break
return np.float32(focal), shift
def recover_focal_shift(points, mask, downsample_size=(64, 64)):
import numpy as np
height, width = points.shape[:2]
uv = normalized_view_plane_uv(width, height, points.dtype)
target_height, target_width = downsample_size
y = np.floor(np.arange(target_height) * height / target_height).astype(int)
x = np.floor(np.arange(target_width) * width / target_width).astype(int)
points_low = points[y[:, None], x[None, :]]
uv_low = uv[y[:, None], x[None, :]]
mask_low = mask[y[:, None], x[None, :]]
points_low = points_low[mask_low]
uv_low = uv_low[mask_low]
if len(points_low) < 2:
return np.float32(1.0), np.float32(0.0)
return solve_focal_shift(uv_low, points_low)
def depth_to_points(depth, intrinsics):
import numpy as np
height, width = depth.shape
u = np.linspace(0.5 / width, 1.0 - 0.5 / width, width, dtype=depth.dtype)
v = np.linspace(0.5 / height, 1.0 - 0.5 / height, height, dtype=depth.dtype)
u, v = np.meshgrid(u, v, indexing="xy")
x = (u - intrinsics[0, 2]) / intrinsics[0, 0] * depth
y = (v - intrinsics[1, 2]) / intrinsics[1, 1] * depth
return np.stack((x, y, depth), axis=-1)
def postprocess(raw_output):
import numpy as np
points = np.asarray(raw_output["points"][0], dtype=np.float32)
normal = np.asarray(raw_output["normal"][0], dtype=np.float32)
mask = np.asarray(raw_output["mask"][0] > 0.5)
metric_scale = np.float32(np.asarray(raw_output["metric_scale"]).reshape(-1)[0])
height, width = mask.shape
aspect_ratio = width / height
focal, shift = recover_focal_shift(points, mask)
diagonal_factor = (1.0 + aspect_ratio**2) ** 0.5
fx = focal / 2.0 * diagonal_factor / aspect_ratio
fy = focal / 2.0 * diagonal_factor
intrinsics = np.array(
((fx, 0.0, 0.5), (0.0, fy, 0.5), (0.0, 0.0, 1.0)),
dtype=np.float32,
)
depth = points[..., 2] + shift
mask &= depth > 0.0
depth = depth * metric_scale
points = depth_to_points(depth, intrinsics)
return {
"points": points.astype(np.float32),
"depth": depth.astype(np.float32),
"normal": normal,
"mask": mask,
"intrinsics": intrinsics,
}
def infer(session, image, resolution_level):
import numpy as np
level = min(max(int(resolution_level), 0), 9)
num_tokens = int(MIN_TOKENS + level / 9.0 * (MAX_TOKENS - MIN_TOKENS))
image_input = np.ascontiguousarray(image.transpose(2, 0, 1)[None], dtype=np.float32)
input_feed = {
"image": image_input,
"num_tokens": np.asarray(num_tokens, dtype=np.int64),
}
output_names = {output.name for output in session.get_outputs()}
scale_name = "metric_scale" if "metric_scale" in output_names else "scale"
names = ("points", "normal", "mask", scale_name)
values = session.run(list(names), input_feed)
raw_output = dict(zip(names, values))
raw_output["metric_scale"] = raw_output.pop(scale_name)
return postprocess(raw_output)

55
src/snapform/onnx_runtime.py Normal file
View File

@@ -0,0 +1,55 @@
"""Shared ONNX Runtime session helpers."""
def select_providers(requested, available):
available = set(available)
providers_for_device = {
"cuda": ("CUDAExecutionProvider", "DmlExecutionProvider"),
"directml": ("DmlExecutionProvider",),
"mps": ("CoreMLExecutionProvider",),
"cpu": ("CPUExecutionProvider",),
}
if requested != "auto":
candidates = providers_for_device[requested]
for provider in candidates:
if provider in available:
return [provider]
if "CPUExecutionProvider" in available:
return ["CPUExecutionProvider"]
raise RuntimeError(
f"{requested.upper()} was requested but the ONNX Runtime "
f"providers {', '.join(candidates)} are not available"
)
priorities = (
"CUDAExecutionProvider",
"CoreMLExecutionProvider",
"DmlExecutionProvider",
"OpenVINOExecutionProvider",
"CPUExecutionProvider",
)
providers = [provider for provider in priorities if provider in available]
if not providers:
raise RuntimeError("ONNX Runtime has no usable execution provider")
return providers
def create_session(path, requested_device="auto"):
import onnxruntime
providers = select_providers(
requested_device,
onnxruntime.get_available_providers(),
)
return onnxruntime.InferenceSession(str(path), providers=providers)
def provider_label(session):
provider = session.get_providers()[0]
labels = {
"CUDAExecutionProvider": "CUDA",
"CoreMLExecutionProvider": "CoreML",
"DmlExecutionProvider": "DirectML",
"OpenVINOExecutionProvider": "OpenVINO",
"CPUExecutionProvider": "CPU",
}
return labels.get(provider, provider)

5
src/snapform/panel.py
View File

@@ -186,7 +186,6 @@ def draw_da3_input_settings(layout, settings, model_property):
layout.prop(settings, "video_frames")
layout.prop(settings, "process_res")
layout.prop(settings, "process_res_method")
layout.prop(settings, "use_ray_pose")
layout.prop(settings, "ref_view_strategy")
@@ -198,7 +197,7 @@ class Da3OutputDebugPanel(bpy.types.Panel):
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 0
bl_order = 1
def draw(self, context):
layout = self.layout
@@ -248,7 +247,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel):
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 1
bl_order = 0
def draw(self, context):
layout = self.layout

4
src/snapform/properties.py
View File

@@ -76,7 +76,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
name="Model",
description="MoGe-2 model used for depth estimation",
items=moge2_model_enum_items,
default=1,
default=0,
)
moge_export_format: bpy.props.EnumProperty(
name="Output Format",
@@ -110,7 +110,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
)
use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Use the slower and more accurate ray-based pose estimation",
description="Legacy task compatibility; ONNX uses the camera decoder",
default=False,
)
ref_view_strategy: bpy.props.EnumProperty(

63
src/snapform/service.py
View File

@@ -74,9 +74,6 @@ class ServiceLock:
class ModelCache:
def __init__(self):
import torch
self.torch = torch
self.ben2_model = None
self.ben2_key = None
self.da3_model = None
@@ -85,75 +82,51 @@ class ModelCache:
self.moge_key = None
def get_ben2(self, model_dir, device):
key = (str(Path(model_dir).resolve()), device.type)
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.ben2_model is not None and self.ben2_key == key:
return self.ben2_model, 0.0
self.release_ben2()
try:
return self._load_ben2(model_dir, device, key)
except self.torch.cuda.OutOfMemoryError:
self.release_da3()
self.release_moge()
return self._load_ben2(model_dir, device, key)
return self._load_ben2(model_dir, requested_device, key)
def _load_ben2(self, model_dir, device, key):
from ben2 import BEN_Base
def _load_ben2(self, model_dir, requested_device, key):
from ben2_onnx import create_session
started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
model = create_session(model_dir, requested_device)
self.ben2_model = model
self.ben2_key = key
elapsed_ms = (time.perf_counter() - started) * 1000.0
return model, elapsed_ms
def get_da3(self, model_dir, device):
key = (str(Path(model_dir).resolve()), device.type)
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.moge_model is not None:
self.release_moge()
if self.da3_model is not None and self.da3_key == key:
return self.da3_model
self.release_da3()
try:
return self._load_da3(model_dir, device, key)
except self.torch.cuda.OutOfMemoryError:
self.release_ben2()
self.release_moge()
return self._load_da3(model_dir, device, key)
return self._load_da3(model_dir, requested_device, key)
def _load_da3(self, model_dir, device, key):
from depth_anything_3.api import DepthAnything3
def _load_da3(self, model_dir, requested_device, key):
from da3_onnx import create_session
model = DepthAnything3.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
model = create_session(model_dir, requested_device)
self.da3_model = model
self.da3_key = key
return model
def get_moge(self, model_dir, device):
key = (str(Path(model_dir).resolve()), device.type)
def get_moge_onnx(self, model_dir, requested_device):
from moge_onnx import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None:
self.release_da3()
if self.moge_model is not None and self.moge_key == key:
return self.moge_model
self.release_moge()
try:
return self._load_moge(model_dir, device, key)
except self.torch.cuda.OutOfMemoryError:
self.release_da3()
self.release_ben2()
return self._load_moge(model_dir, device, key)
def _load_moge(self, model_dir, device, key):
worker.prepare_moge_runtime()
from moge.model.v2 import MoGeModel
checkpoint = Path(model_dir) / "model.pt"
model = MoGeModel.from_pretrained(checkpoint)
model = model.to(device=device)
model.eval()
model = create_session(model_dir, requested_device)
self.moge_model = model
self.moge_key = key
return model
@@ -201,8 +174,6 @@ class ModelCache:
def _release_cuda_cache(self):
gc.collect()
if self.torch.cuda.is_available():
self.torch.cuda.empty_cache()
class ServiceState:

378
src/snapform/worker.py
View File

@@ -6,11 +6,17 @@ import shutil
import tempfile
import threading
import time
import sys
from pathlib import Path
from urllib.parse import quote
from urllib.request import Request, urlopen
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = {
".mp4",
@@ -434,6 +440,92 @@ def export_prediction_attributes(prediction, output_directory):
return output_path
def export_da3_mini_npz(prediction, output_directory):
import numpy as np
values = {"depth": np.round(prediction.depth, 8)}
for name in ("conf", "extrinsics", "intrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.round(value, 2) if name == "conf" else value
output_path = Path(output_directory) / "results.npz"
np.savez_compressed(output_path, **values)
return output_path
def export_da3_depth_preview(prediction, output_directory):
import numpy as np
from PIL import Image
target = Path(output_directory) / "depth_vis"
target.mkdir(parents=True, exist_ok=True)
for index, depth in enumerate(prediction.depth):
valid = np.isfinite(depth) & (depth > 0.0)
normalized = np.zeros(depth.shape, dtype=np.float32)
if valid.any():
low, high = np.percentile(depth[valid], (2.0, 98.0))
normalized[valid] = np.clip(
(depth[valid] - low) / max(float(high - low), 1e-6),
0.0,
1.0,
)
color = np.stack(
(
normalized,
1.0 - np.abs(normalized * 2.0 - 1.0),
1.0 - normalized,
),
axis=-1,
)
color = np.asarray(color * 255.0, dtype=np.uint8)
source = np.asarray(prediction.processed_images[index], dtype=np.uint8)
combined = np.concatenate((source, color), axis=1)
Image.fromarray(combined, mode="RGB").save(
target / f"{index:04d}.jpg",
quality=95,
)
def export_da3_point_cloud(
prediction,
output_directory,
confidence_percentile=40.0,
num_max_points=1_000_000,
):
import numpy as np
import trimesh
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
threshold = confidence_percentile_threshold(
depth,
confidence,
confidence_percentile,
)
point_sets = []
color_sets = []
for index in range(depth.shape[0]):
points = unproject_depth_frame(
depth[index],
prediction.intrinsics[index],
prediction.extrinsics[index],
)
valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0)
if threshold is not None:
valid &= confidence[index] >= threshold
point_sets.append(points[valid])
color_sets.append(prediction.processed_images[index][valid])
points = np.concatenate(point_sets, axis=0)
colors = np.concatenate(color_sets, axis=0)
maximum = max(int(num_max_points), 1)
if len(points) > maximum:
indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64)
points = points[indices]
colors = colors[indices]
scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors))
scene.export(Path(output_directory) / SCENE_FILENAME)
def remove_background_faces(faces, foreground_mask):
"""Return only triangles whose vertices all belong to the foreground."""
foreground_points = foreground_mask.reshape(-1)
@@ -1133,29 +1225,10 @@ def format_rate(rate):
return f"{size:.1f} {units[unit_index]}"
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")
if requested == "mps":
if not torch.backends.mps.is_available():
raise RuntimeError("MPS was requested but is not available")
return torch.device("mps")
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def infer_foreground(
input_path,
model_dir,
device,
torch,
model_cache=None,
cancel_check=None,
):
@@ -1165,14 +1238,13 @@ def infer_foreground(
check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
owns_model = model_cache is None
if owns_model:
from ben2 import BEN_Base
from ben2_onnx import create_session, infer
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
if owns_model:
load_started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
model = create_session(model_dir, requested_device)
load_ms = (time.perf_counter() - load_started) * 1000.0
else:
model, load_ms = model_cache.get_ben2(model_dir, device)
@@ -1180,22 +1252,10 @@ def infer_foreground(
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
if device.type == "cuda":
torch.cuda.synchronize()
inference_started = time.perf_counter()
with torch.inference_mode():
foreground = model.inference(
image,
refine_foreground=False,
)
foreground = infer(model, image)
foreground = preserve_source_rgb(image, foreground)
if device.type == "cuda":
torch.cuda.synchronize()
inference_ms = (time.perf_counter() - inference_started) * 1000.0
if owns_model:
del model
if device.type == "cuda":
torch.cuda.empty_cache()
check_cancelled(cancel_check)
return foreground, load_ms, inference_ms
@@ -1216,7 +1276,6 @@ def infer_foreground_frames(
images,
model_dir,
device,
torch,
output_directory,
status_path,
start_progress,
@@ -1238,7 +1297,6 @@ def infer_foreground_frames(
image_path,
model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
@@ -1260,14 +1318,9 @@ def remove_background(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check)
if model_cache is None:
import torch
else:
torch = model_cache.torch
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True)
device = select_device(args.device, torch)
device = args.device
video_frames = getattr(args, "video_frames", 16)
start_frame = getattr(args, "start_frame", 1)
@@ -1279,14 +1332,13 @@ def remove_background(args, model_cache=None, cancel_check=None):
cancel_check=cancel_check,
)
try:
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
write_status(args.status, 0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
write_status(args.status, 0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
args.model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
@@ -1305,7 +1357,6 @@ def remove_background(args, model_cache=None, cancel_check=None):
images,
args.model_dir,
device,
torch,
output_directory,
args.status,
0.12,
@@ -1517,41 +1568,33 @@ def export_moge_debug_mesh(
export_format,
):
import numpy as np
import utils3d
from moge.utils.io import save_glb, save_ply
import trimesh
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
uv = utils3d.np.uv_map(*depth.shape)
result = utils3d.np.build_mesh_from_map(
clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04)
faces = mesh_triangle_faces(
points,
image,
uv,
normal,
mask=clean_mask,
tri=True,
depth,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=clean_mask,
)
faces, vertices, colors, vertex_uvs, vertex_normals = result
vertices = points.reshape(-1, 3)
colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8)
vertex_normals = normal.reshape(-1, 3)
vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32)
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
if export_format == "glb":
save_glb(
output_directory / "mesh.glb",
vertices,
faces,
vertex_uvs,
(image * 255.0).astype(np.uint8),
vertex_normals,
)
return
if export_format == "ply":
save_ply(
output_directory / "pointcloud.ply",
vertices,
np.zeros((0, 3), dtype=np.int32),
colors,
vertex_normals,
)
mesh = trimesh.Trimesh(
vertices=vertices,
faces=faces,
vertex_colors=colors,
vertex_normals=vertex_normals,
process=False,
)
mesh.remove_unreferenced_vertices()
filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply"
mesh.export(Path(output_directory) / filename)
def export_moge_surface_archive(
@@ -1564,7 +1607,6 @@ def export_moge_surface_archive(
foreground_mask=None,
):
import numpy as np
import utils3d
height, width = depth.shape
finite_mask = np.isfinite(depth) & (depth > 0.0)
@@ -1572,34 +1614,28 @@ def export_moge_surface_archive(
if foreground_mask is not None:
finite_mask &= foreground_mask
depth_discontinuity = depth_discontinuity_ratio(depth)
result = utils3d.np.build_mesh_from_map(
faces = mesh_triangle_faces(
points,
utils3d.np.uv_map(height, width),
normal,
mask.astype(np.float32),
depth_discontinuity.astype(np.float32),
depth,
mask=finite_mask,
tri=True,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=finite_mask,
)
(
faces,
vertices,
vertex_uvs,
_vertex_normals,
vertex_mask,
vertex_depth_discontinuity,
point_depth,
) = result
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
vertices = points.reshape(-1, 3)
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
uu, vv = np.meshgrid(u, v)
vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2)
pixel_intrinsics = intrinsics.copy()
pixel_intrinsics[0] *= width
pixel_intrinsics[1] *= height
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = point_depth.astype(np.float32)
face_mask = vertex_mask[faces].all(axis=1)
face_depth_discontinuity = vertex_depth_discontinuity[faces].max(axis=1)
masked_depth = point_depth[vertex_mask.astype(bool)]
point_depth = depth.reshape(-1).astype(np.float32)
face_mask = mask.reshape(-1)[faces].all(axis=1)
face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
masked_depth = point_depth[mask.reshape(-1).astype(bool)]
reference_depth = np.median(masked_depth if len(masked_depth) else point_depth)
np.savez(
output_path,
@@ -1619,14 +1655,13 @@ def export_moge_surface_archive(
def run_moge(args, model_cache=None, cancel_check=None):
import numpy as np
from PIL import Image
from moge_onnx import (
create_session,
infer as infer_moge_onnx,
provider_label,
)
prepare_moge_runtime()
write_status(args.status, 0.02, "Starting MoGe-2")
if model_cache is None:
import torch
from moge.model.v2 import MoGeModel
else:
torch = model_cache.torch
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
@@ -1637,15 +1672,15 @@ def run_moge(args, model_cache=None, cancel_check=None):
cancel_check=cancel_check,
)
try:
device = select_device(args.device, torch)
write_status(args.status, 0.08, f"Loading MoGe-2 on {device.type}")
if model_cache is None:
checkpoint = Path(args.model_dir) / "model.pt"
model = MoGeModel.from_pretrained(checkpoint)
model = model.to(device=device)
model.eval()
model = create_session(args.model_dir, args.device)
else:
model = model_cache.get_moge(args.model_dir, device)
model = model_cache.get_moge_onnx(args.model_dir, args.device)
write_status(
args.status,
0.08,
f"Loading MoGe-2 on {provider_label(model)}",
)
depths = []
normals = []
@@ -1662,21 +1697,17 @@ def run_moge(args, model_cache=None, cancel_check=None):
)
with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device)
source_images.append(image)
with torch.inference_mode():
prediction = model.infer(
image_tensor,
resolution_level=int(args.resolution_level),
force_projection=True,
apply_mask=False,
use_fp16=device.type == "cuda",
)
depth = prediction["depth"].detach().float().cpu().numpy()
normal = prediction["normal"].detach().float().cpu().numpy()
mask = prediction["mask"].detach().cpu().numpy().astype(bool)
point = prediction["points"].detach().float().cpu().numpy()
intrinsic = prediction["intrinsics"].detach().float().cpu().numpy()
prediction = infer_moge_onnx(
model,
image,
int(args.resolution_level),
)
depth = prediction["depth"]
normal = prediction["normal"]
mask = prediction["mask"].astype(bool)
point = prediction["points"]
intrinsic = prediction["intrinsics"]
depths.append(depth)
normals.append(normal)
masks.append(mask)
@@ -1727,8 +1758,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
foreground, _, _ = infer_foreground(
images[0],
remove_background_model_dir,
device,
torch,
args.device,
model_cache=model_cache,
cancel_check=cancel_check,
)
@@ -1743,8 +1773,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
infer_foreground_frames(
images,
remove_background_model_dir,
device,
torch,
args.device,
output_directory,
args.status,
0.87,
@@ -1781,21 +1810,11 @@ def run_moge(args, model_cache=None, cancel_check=None):
shutil.rmtree(temporary_directory, ignore_errors=True)
def prepare_moge_runtime():
"""Use PyTorch SDPA because xFormers does not support every CUDA GPU."""
os.environ["XFORMERS_DISABLED"] = "1"
def reconstruct(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check)
if model_cache is None:
import torch
from depth_anything_3.api import DepthAnything3
else:
torch = model_cache.torch
from depth_anything_3.utils.export import export
from da3_onnx import create_session, infer
from onnx_runtime import provider_label
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=True)
@@ -1810,7 +1829,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
try:
check_cancelled(cancel_check)
device = select_device(args.device, torch)
remove_background_model_dir = getattr(
args,
"remove_background_model_dir",
@@ -1823,8 +1841,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
foreground, load_ms, inference_ms = infer_foreground(
images[0],
remove_background_model_dir,
device,
torch,
args.device,
model_cache=model_cache,
cancel_check=cancel_check,
)
@@ -1841,8 +1858,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
load_ms, inference_ms = infer_foreground_frames(
images,
remove_background_model_dir,
device,
torch,
args.device,
output_directory,
args.status,
0.12,
@@ -1858,13 +1874,16 @@ def reconstruct(args, model_cache=None, cancel_check=None):
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)
write_status(args.status, 0.15, f"Loading model on {device.type}")
write_status(args.status, 0.15, "Loading DA3 ONNX model")
if model_cache is None:
model = DepthAnything3.from_pretrained(args.model_dir)
model = model.to(device=device)
model.eval()
model = create_session(args.model_dir, args.device)
else:
model = model_cache.get_da3(args.model_dir, device)
model = model_cache.get_da3(args.model_dir, args.device)
write_status(
args.status,
0.20,
f"Loaded DA3 on {provider_label(model)}",
)
check_cancelled(cancel_check)
write_status(
@@ -1872,27 +1891,14 @@ def reconstruct(args, model_cache=None, cancel_check=None):
0.25,
f"Estimating depth for {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,
)
prediction = infer(
model,
[str(path) for path in images],
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,
)
check_cancelled(cancel_check)
write_status(args.status, 0.88, f"Exporting {args.export_format}")
if args.export_format == "mesh":
@@ -1921,12 +1927,20 @@ def reconstruct(args, model_cache=None, cancel_check=None):
),
progress_callback=mesh_progress,
)
else:
export(
elif args.export_format == "mini_npz":
export_da3_mini_npz(prediction, output_directory)
elif args.export_format == "depth_vis":
export_da3_depth_preview(prediction, output_directory)
elif args.export_format == "glb":
export_da3_point_cloud(
prediction,
args.export_format,
str(output_directory),
**export_kwargs,
output_directory,
args.conf_thresh_percentile,
args.num_max_points,
)
else:
raise ValueError(
f"DA3 ONNX does not support export format: {args.export_format}"
)
if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
write_status(
@@ -1989,7 +2003,7 @@ def parser():
background.add_argument("--status", required=True)
background.add_argument(
"--device",
choices=("auto", "cuda", "mps", "cpu"),
choices=("auto", "cuda", "directml", "mps", "cpu"),
required=True,
)
background.add_argument("--video-frames", type=int, default=16)
@@ -2040,7 +2054,7 @@ def parser():
reconstruction.add_argument("--start-frame", type=int, default=1)
reconstruction.add_argument(
"--device",
choices=("auto", "cuda", "mps", "cpu"),
choices=("auto", "cuda", "directml", "mps", "cpu"),
required=True,
)
reconstruction.add_argument(