nan
1c98b21235
- Introduced ModelSpec dataclass to define model attributes. - Added multiple models including DA3 Small, Base, Giant 1.1, Nested Giant Large 1.1, and MoGe-2 variants. - Implemented functions to retrieve model details and metadata. - Included BEN2 model specifications for background removal.
233 lines
8.5 KiB
Python
233 lines
8.5 KiB
Python
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"],
|
|
"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)
|