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

Add ONNX model inference implementations for MoGe-2, BEN2, and DA3

Browse Source 
- 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:
马南
2026-08-17 10:37:55 +08:00
parent b3a07cde32
commit c5b77c5c85
58 changed files with 5840 additions and 9931 deletions
Download Patch File Download Diff File Expand all files Collapse all files

388
src/snapform/blender/depth_plane.py
View File

@@ -1,5 +1,6 @@
import re
import shutil
import tempfile
from pathlib import Path
import bpy
@@ -53,7 +54,7 @@ VIDEO_NAME_SUFFIXES = (
def image_depth_model_items(owner, context):
from ..models import depth_plane_model_enum_items
from ..properties import depth_plane_model_enum_items
return depth_plane_model_enum_items(owner, context)
@@ -218,19 +219,17 @@ def process_resolution_scale_property():
def draw_process_resolution(layout, owner, image):
process_res = scaled_process_resolution(
image,
owner.snapform_process_res_factor,
owner.process_res_factor,
)
layout.prop(
owner,
"snapform_process_res_factor",
"process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
def save_packed_image(image, scene):
from ..client import temporary_input_path
output_path = temporary_input_path()
output_path = Path(tempfile.mkdtemp(prefix="snapform-input-")) / "input.png"
image_settings = scene.render.image_settings
original_file_format = image_settings.file_format
try:
@@ -286,9 +285,7 @@ def prepare_image_input(image, scene):
return source_path, False
sequence_frames = collect_sequence_frames(image)
if sequence_frames:
from ..client import temporary_input_directory
temporary_directory = Path(temporary_input_directory())
temporary_directory = Path(tempfile.mkdtemp(prefix="snapform-input-"))
try:
for index, frame_path in enumerate(sequence_frames):
suffix = frame_path.suffix.lower() or ".jpg"
@@ -1868,90 +1865,6 @@ def build_depth_plane(
)
class BuildDepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_depth_plane"
bl_label = "Build Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_data_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_color_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
npz_path = Path(self.snapform_npz_path)
if not npz_path.is_file():
self.report(
{"ERROR"},
f"Depth result does not exist: {npz_path}",
)
return {"CANCELLED"}
try:
foreground_image = None
if self.snapform_background_path:
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
color_image = None
# The foreground result already carries the color, so the original
# frame sequence is unused when background removal is available.
if self.snapform_color_frames_dir and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
build_depth_plane(
context,
source_object,
npz_path,
self.snapform_delete_source,
unshaded=self.snapform_unshaded,
depth_name=self.snapform_depth_data_name,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
def build_moge2_depth_plane(
context,
source_object,
@@ -2048,177 +1961,6 @@ def build_moge2_depth_plane(
)
class BuildMoge2DepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_moge2_depth_plane"
bl_label = "Build MoGe-2 Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(options={"HIDDEN"})
snapform_npz_path: bpy.props.StringProperty(options={"HIDDEN"})
snapform_background_path: bpy.props.StringProperty(options={"HIDDEN"})
snapform_color_frames_dir: bpy.props.StringProperty(options={"HIDDEN"})
snapform_data_frames_dir: bpy.props.StringProperty(options={"HIDDEN"})
snapform_delete_source: bpy.props.BoolProperty(default=True, options={"HIDDEN"})
snapform_unshaded: bpy.props.BoolProperty(default=False, options={"HIDDEN"})
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report({"ERROR"}, "The source image empty is no longer available")
return {"CANCELLED"}
try:
foreground_image = None
if self.snapform_background_path:
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
color_image = None
# The foreground result already carries the color, so the original
# frame sequence is unused when background removal is available.
if self.snapform_color_frames_dir and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
build_moge2_depth_plane(
context,
source_object,
Path(self.snapform_npz_path),
self.snapform_delete_source,
unshaded=self.snapform_unshaded,
color_image=color_image,
foreground_image=foreground_image,
data_frames_directory=(
self.snapform_data_frames_dir or None
),
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class BuildOriginalDepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_original_depth_plane"
bl_label = "Build Original Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_color_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
try:
foreground_image = None
if self.snapform_background_path:
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
color_image = None
# The foreground result already carries the color, so the original
# frame sequence is unused when background removal is available.
if self.snapform_color_frames_dir and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
depth_image = None
if self.snapform_depth_npz_path:
depth_arrays = extract_depth_arrays(
self.snapform_depth_npz_path
)
if len(depth_arrays) == 1:
depth_image = create_depth_image(
depth_data_name(source_object),
depth_arrays[0],
)
else:
depth_image = create_depth_sequence_image(
depth_data_name(source_object),
depth_arrays,
context.scene,
self.snapform_depth_frames_dir,
)
build_image_depth_plane(
context,
source_object,
self.snapform_delete_source,
unshaded=self.snapform_unshaded,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
depth_image=depth_image,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class ApplyRemovedBackground(bpy.types.Operator):
bl_idname = "snapform.apply_removed_background"
bl_label = "Apply Removed Background"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_output_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, _context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
try:
create_removed_background_image(
source_object,
self.snapform_output_path,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class RemoveImageBackground(bpy.types.Operator):
bl_idname = "snapform.remove_image_background"
bl_label = "Remove Background"
@@ -2227,15 +1969,13 @@ class RemoveImageBackground(bpy.types.Operator):
@classmethod
def poll(cls, context):
from ..operators import active_job
from ..client import service_busy, model_ready
from ..models import BEN2_MODEL
from ..job_runtime import runtime
from ..properties import BEN2_MODEL_KEY, model_ready
return (
is_image_empty(context.object)
and active_job() is None
and not service_busy()
and model_ready(BEN2_MODEL.key)
and not runtime.server_busy()
and model_ready(BEN2_MODEL_KEY)
)
def execute(self, context):
@@ -2255,9 +1995,9 @@ class RemoveImageBackground(bpy.types.Operator):
return {"CANCELLED"}
result = bpy.ops.snapform.process_image_background(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_delete_input=delete_input,
input_path=str(source_path),
source_object_name=source_object.name,
delete_input=delete_input,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2266,42 +2006,42 @@ class RemoveImageBackground(bpy.types.Operator):
return {"CANCELLED"}
class ConvertImageToDepthPlane(bpy.types.Operator):
bl_idname = "snapform.convert_image_to_depth_plane"
class ConvertToDepthPlane(bpy.types.Operator):
bl_idname = "snapform.convert_to_depth_plane"
bl_label = "Convert to Depth Plane"
bl_description = "Convert an image to a displaced depth plane"
bl_options = {"REGISTER"}
snapform_unshaded: bpy.props.BoolProperty(
unshaded: bpy.props.BoolProperty(
name="Unshaded",
description="Display the source color through an Emission shader",
default=False,
)
snapform_remove_background: bpy.props.BoolProperty(
remove_background: bpy.props.BoolProperty(
name="Remove Background",
description="Remove the background during conversion with BEN2",
default=False,
)
snapform_model: bpy.props.EnumProperty(
model: bpy.props.EnumProperty(
name="Depth From",
description="Model or original image used to build depth",
items=image_depth_model_items,
default=0,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
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(
process_res_factor: process_resolution_scale_property()
resolution_level: bpy.props.IntProperty(
name="Resolution Level",
description="MoGe-2 inference resolution level; higher is finer and slower",
default=5,
min=0,
max=9,
)
snapform_keep_original: bpy.props.BoolProperty(
keep_original: bpy.props.BoolProperty(
name="Keep Original",
description="Keep the source Image Empty after conversion",
default=False,
@@ -2309,13 +2049,11 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
@classmethod
def poll(cls, context):
from ..operators import active_job
from ..client import service_busy
from ..job_runtime import runtime
return (
is_image_empty(context.object)
and active_job() is None
and not service_busy()
and not runtime.server_busy()
)
def invoke(self, context, _event):
@@ -2323,15 +2061,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
def draw(self, context):
layout = self.layout
layout.prop(self, "snapform_model")
if self.snapform_model in {"SMALL", "BASE"}:
layout.prop(self, "model")
if self.model in {"SMALL", "BASE"}:
draw_process_resolution(layout, self, context.object.data)
elif self.snapform_model.startswith("MOGE2_"):
layout.prop(self, "snapform_moge_resolution_level")
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_unshaded")
elif self.model.startswith("MOGE2_"):
layout.prop(self, "resolution_level")
layout.prop(self, "remove_background")
layout.prop(self, "unshaded")
layout.separator(type="LINE")
layout.prop(self, "snapform_keep_original")
layout.prop(self, "keep_original")
def execute(self, context):
source_object = context.object
@@ -2339,7 +2077,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
context.scene.snapform_settings.video_frames = (
default_extract_frames(source_object)
)
if self.snapform_model == "ORIGINAL":
if self.model == "ORIGINAL":
try:
source_path, delete_input = prepare_image_input(
source_object.data,
@@ -2350,15 +2088,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
run_original = (
bpy.ops.snapform.generate_original_background
if self.snapform_remove_background
if self.remove_background
else bpy.ops.snapform.extract_original_frames
)
result = run_original(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
input_path=str(source_path),
source_object_name=source_object.name,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2366,7 +2104,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
remove_temporary_input(source_path)
return {"CANCELLED"}
if self.snapform_model.startswith("MOGE2_"):
if self.model.startswith("MOGE2_"):
try:
source_path, delete_input = prepare_image_input(
source_object.data,
@@ -2376,16 +2114,14 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.generate_moge2_depth(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_moge_resolution_level=(
self.snapform_moge_resolution_level
),
snapform_model_key=self.snapform_model,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
snapform_remove_background=self.snapform_remove_background,
input_path=str(source_path),
source_object_name=source_object.name,
resolution_level=self.resolution_level,
model=self.model,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
remove_background=self.remove_background,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2397,7 +2133,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
try:
process_res = scaled_process_resolution(
image,
self.snapform_process_res_factor,
self.process_res_factor,
)
except ValueError as error:
self.report({"ERROR"}, str(error))
@@ -2412,16 +2148,16 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
result = bpy.ops.snapform.generate_da3_depth(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_depth_data_name=depth_data_name(source_object),
snapform_model_key=self.snapform_model,
snapform_use_ray_pose=self.snapform_use_ray_pose,
snapform_process_res=process_res,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
snapform_remove_background=self.snapform_remove_background,
input_path=str(source_path),
source_object_name=source_object.name,
depth_data_name=depth_data_name(source_object),
model=self.model,
use_ray_pose=self.use_ray_pose,
process_res=process_res,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
remove_background=self.remove_background,
)
if "CANCELLED" in result:
if delete_input:
@@ -2435,11 +2171,11 @@ class SnapFormImageMenu(bpy.types.Menu):
bl_label = "SnapForm"
def draw(self, _context):
from ..client import env_installed
from ..job_runtime import runtime
layout = self.layout
layout.operator_context = "INVOKE_DEFAULT"
if not env_installed():
if not runtime.environment_ready():
layout.operator(
"snapform.install_environment",
text="Install Environment",
@@ -2447,12 +2183,12 @@ class SnapFormImageMenu(bpy.types.Menu):
)
return
layout.operator(
"snapform.convert_image_to_depth_plane",
"snapform.convert_to_depth_plane",
text="Convert to Depth Plane",
icon="MOD_DISPLACE",
)
layout.operator(
"snapform.convert_image_to_mesh",
"snapform.convert_to_surface",
text="Convert to Surface",
icon="MESH_GRID",
)