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)