Add ONNX model inference implementations for MoGe-2, BEN2, and DA3
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing. - Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing. - Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Developed unit tests for server functionalities in `test_server.py`, including model caching and job management. - Enhanced job server to handle frame extraction without requiring models.
This commit is contained in:
233
src/snapform/server/da3.py
Normal file
233
src/snapform/server/da3.py
Normal file
@@ -0,0 +1,233 @@
|
||||
import shutil
|
||||
|
||||
from .background import infer_foreground, infer_foreground_frames
|
||||
from .geometry import (
|
||||
export_da3_depth_preview,
|
||||
export_da3_mini_npz,
|
||||
export_da3_point_cloud,
|
||||
export_image_mesh_data,
|
||||
export_prediction_attributes,
|
||||
export_textured_mesh_glb,
|
||||
)
|
||||
from .media import (
|
||||
copy_frames,
|
||||
flatten_exported_paths,
|
||||
prepare_input,
|
||||
wait_for_exported_paths,
|
||||
)
|
||||
|
||||
|
||||
def merge_da3_frame_predictions(predictions):
|
||||
"""Combine independent single-frame DA3 results for existing exporters."""
|
||||
import numpy as np
|
||||
|
||||
if not predictions:
|
||||
raise ValueError("DA3 returned no frame predictions")
|
||||
merged = predictions[0]
|
||||
for name in (
|
||||
"depth",
|
||||
"conf",
|
||||
"extrinsics",
|
||||
"intrinsics",
|
||||
"processed_images",
|
||||
):
|
||||
values = [getattr(prediction, name) for prediction in predictions]
|
||||
setattr(merged, name, np.concatenate(values, axis=0))
|
||||
return merged
|
||||
|
||||
|
||||
def generate_da3(context, parameters):
|
||||
"""Run DA3 inference and return the artifact used by the requesting Job."""
|
||||
context.progress(0.02, "Starting Environment")
|
||||
context.check_cancelled()
|
||||
from .onnx_da3 import infer
|
||||
from .onnx_runtime import provider_label
|
||||
|
||||
output_directory = context.directory
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
context.progress( 0.08, "Preparing input")
|
||||
model_manager = context.resource("model_manager")
|
||||
model_directory = model_manager.directory(parameters["model"])
|
||||
device = parameters["device"]
|
||||
export_format = parameters["export_format"]
|
||||
images, temporary_directory = prepare_input(
|
||||
parameters["input"],
|
||||
parameters.get("video_frames", 16),
|
||||
output_directory,
|
||||
start_frame=parameters.get("start_frame", 1),
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
|
||||
try:
|
||||
context.check_cancelled()
|
||||
remove_background_model = parameters.get("remove_background_model", "")
|
||||
remove_background_model_dir = (
|
||||
model_manager.directory(remove_background_model)
|
||||
if remove_background_model
|
||||
else ""
|
||||
)
|
||||
foreground_images = []
|
||||
if remove_background_model_dir:
|
||||
context.progress( 0.12, "Removing image background")
|
||||
if len(images) == 1:
|
||||
foreground, load_ms, inference_ms = infer_foreground(
|
||||
images[0],
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
context.check_cancelled()
|
||||
foreground.convert("RGBA").save(
|
||||
output_directory / "foreground.png",
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
if not (output_directory / "foreground.png").is_file():
|
||||
raise RuntimeError("BEN2 did not write a foreground image")
|
||||
foreground_images = [output_directory / "foreground.png"]
|
||||
else:
|
||||
load_ms, inference_ms = infer_foreground_frames(
|
||||
images,
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
output_directory,
|
||||
context.progress,
|
||||
0.12,
|
||||
0.05,
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
foreground_images = sorted(
|
||||
output_directory.glob("foreground.*.png")
|
||||
)
|
||||
if remove_background_model_dir:
|
||||
print(
|
||||
f"BEN2 inference {inference_ms:.0f} ms; "
|
||||
f"model load {load_ms:.0f} ms"
|
||||
)
|
||||
context.progress( 0.15, "Loading DA3 ONNX model")
|
||||
model = model_manager.get_da3(model_directory, device)
|
||||
context.progress(
|
||||
0.20,
|
||||
f"Loaded DA3 on {provider_label(model)}",
|
||||
)
|
||||
|
||||
context.check_cancelled()
|
||||
context.progress(
|
||||
0.25,
|
||||
f"Estimating depth for {len(images)} frame(s)",
|
||||
)
|
||||
inference_options = {
|
||||
"process_res": parameters["process_res"],
|
||||
"process_res_method": parameters["process_res_method"],
|
||||
"use_ray_pose": bool(parameters["use_ray_pose"]),
|
||||
"ref_view_strategy": parameters["ref_view_strategy"],
|
||||
}
|
||||
try:
|
||||
if parameters.get("inference_mode", "joint") == "per_frame":
|
||||
predictions = []
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
context.progress(
|
||||
0.25 + 0.60 * index / len(images),
|
||||
f"Estimating DA3 depth ({index + 1}/{len(images)})",
|
||||
)
|
||||
predictions.append(
|
||||
infer(model, [str(image)], **inference_options)
|
||||
)
|
||||
prediction = merge_da3_frame_predictions(predictions)
|
||||
else:
|
||||
prediction = infer(
|
||||
model,
|
||||
[str(path) for path in images],
|
||||
**inference_options,
|
||||
)
|
||||
except UnicodeDecodeError as error:
|
||||
# DirectML reports failures with a locale-specific message that
|
||||
# ONNX Runtime cannot always decode, masking the real error.
|
||||
if provider_label(model) == "DirectML":
|
||||
raise RuntimeError(
|
||||
"DA3 inference failed on DirectML. Switch the device to "
|
||||
"CPU in SnapForm Preferences and retry. "
|
||||
f"Original error: {error}"
|
||||
) from error
|
||||
raise
|
||||
context.check_cancelled()
|
||||
context.progress( 0.88, f"Exporting {export_format}")
|
||||
if export_format == "mesh":
|
||||
export_textured_mesh_glb(
|
||||
prediction,
|
||||
images,
|
||||
output_directory,
|
||||
parameters.get("conf_thresh_percentile", 40.0),
|
||||
bool(parameters.get("mesh_filter_edges", True)),
|
||||
float(parameters.get("mesh_max_edge_length", 0.1)),
|
||||
)
|
||||
elif export_format == "mesh_data":
|
||||
def mesh_progress(step, total, label):
|
||||
context.progress(
|
||||
0.88 + 0.09 * step / total,
|
||||
f"Preparing surface data {step}/{total}: {label}",
|
||||
)
|
||||
|
||||
export_image_mesh_data(
|
||||
prediction,
|
||||
images[0],
|
||||
output_directory,
|
||||
foreground_path=(
|
||||
foreground_images[0] if foreground_images else None
|
||||
),
|
||||
progress_callback=mesh_progress,
|
||||
)
|
||||
elif export_format == "mini_npz":
|
||||
export_da3_mini_npz(prediction, output_directory)
|
||||
elif export_format == "depth_vis":
|
||||
export_da3_depth_preview(prediction, output_directory)
|
||||
elif export_format == "glb":
|
||||
export_da3_point_cloud(
|
||||
prediction,
|
||||
output_directory,
|
||||
parameters["conf_thresh_percentile"],
|
||||
parameters["num_max_points"],
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"DA3 ONNX does not support export format: {export_format}"
|
||||
)
|
||||
if export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
|
||||
context.progress(
|
||||
0.97,
|
||||
"Writing structured debug attributes",
|
||||
)
|
||||
export_prediction_attributes(prediction, output_directory)
|
||||
context.check_cancelled()
|
||||
|
||||
result_paths = wait_for_exported_paths(
|
||||
output_directory,
|
||||
export_format,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
if not result_paths:
|
||||
raise RuntimeError(
|
||||
f"No {export_format} files were exported"
|
||||
)
|
||||
result_paths = flatten_exported_paths(
|
||||
output_directory,
|
||||
result_paths,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
context.check_cancelled()
|
||||
copy_frames(images, output_directory)
|
||||
context.progress(
|
||||
1.0,
|
||||
f"Exported {len(result_paths)} {export_format} file(s)",
|
||||
)
|
||||
if context.job_type == "generate-da3-surface":
|
||||
return {"surface": "mesh.npz"}
|
||||
if context.job_type == "generate-da3-depth":
|
||||
return {"depth": "results.npz"}
|
||||
return {"directory": "."}
|
||||
finally:
|
||||
if temporary_directory is not None:
|
||||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||
Reference in New Issue
Block a user