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

feat: Add MoGe-2 surface generation and export functionality

Browse Source 
- Introduced new enum items for depth plane and surface models with caching.
- Added MOGE2_MESH_JOB to operators for converting images to surfaces.
- Implemented surface mesh building and exporting in RunJob operator.
- Enhanced UI to include export format options for MoGe-2.
- Created validation tools for MoGe-2 surface assets and depth planes.
- Updated geometry nodes to support new surface generation features.
- Refactored existing tests to accommodate new functionality and ensure stability.
This commit is contained in:
马南
2026-08-13 22:07:36 +08:00
parent a6370d678d
commit caeb6d706c
17 changed files with 902 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
src/snapform/assets/O_SnapForm.blend LFS
View File

Binary file not shown.

5
src/snapform/client.py
View File

@@ -38,6 +38,7 @@ JOB_FOLDER_NAMES = {
"RECONSTRUCT": "export-reconstruction",
"IMAGE_DEPTH": "generate-depth-plane",
"MOGE2_DEPTH": "generate-moge2-depth-plane",
"MOGE2_MESH": "generate-moge2-surface",
"MOGE2_DEBUG": "test-moge2-api",
"IMAGE_MESH": "generate-image-mesh",
"REMOVE_BACKGROUND": "remove-background",
@@ -909,6 +910,8 @@ def moge_request(
start_frame=1,
device=None,
model_key=None,
debug_export_format="",
surface_mesh=False,
):
from .models import BEN2_MODEL, MOGE2_MODEL
@@ -929,6 +932,8 @@ def moge_request(
else resolution_level
),
"include_points": bool(include_points),
"debug_export_format": debug_export_format,
"surface_mesh": bool(surface_mesh),
"video_frames": settings.video_frames,
"start_frame": start_frame,
"device": (device or settings.device).lower(),

45
src/snapform/image_depth.py
View File

@@ -19,6 +19,7 @@ MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement"
BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth"
DEPTH_SCALE_ATTRIBUTE = "o_depth_scale"
SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer"
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Surface Layer"
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer"
MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer"
DA3_PATCH_SIZE = 14
@@ -57,9 +58,6 @@ def image_depth_model_items(owner, context):
return depth_plane_model_enum_items(owner, context)
_last_depth_plane_model = "MOGE2_VITB_NORMAL"
def is_image_empty(candidate):
return bool(
candidate
@@ -1006,10 +1004,13 @@ def create_moge_material(
uv.location = (-900.0, 200.0)
color = nodes.new("ShaderNodeTexImage")
color.image = color_image
color.interpolation = "Linear"
color.extension = "EXTEND"
color.location = (-650.0, 300.0)
depth = nodes.new("ShaderNodeTexImage")
depth.image = depth_image
depth.interpolation = "Linear"
depth.extension = "EXTEND"
depth.location = (-650.0, -100.0)
if frame_count > 1:
configure_image_user(
@@ -1127,11 +1128,11 @@ def create_plane_mesh(source_object, bounds=None, subdivision_level=5):
return mesh
def material_node_group(depth_plane=False):
def material_node_group(depth_plane=False, surface_group_name=None):
group_name = (
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME
if depth_plane
else SURFACE_MATERIAL_NODE_GROUP_NAME
else surface_group_name or SURFACE_MATERIAL_NODE_GROUP_NAME
)
existing = bpy.data.node_groups.get(group_name)
if existing is not None:
@@ -1269,7 +1270,7 @@ def material_node_group(depth_plane=False):
group.description = (
"Prepare DA3 depth-plane material outputs."
if depth_plane
else "Prepare DA3 surface material outputs."
else f"Prepare {group_name} material outputs."
)
group.use_fake_user = True
return group
@@ -1293,6 +1294,7 @@ def create_material(
depth_frame_offset=0,
depth_frame_duration=None,
use_cyclic=None,
surface_group_name=None,
):
material = bpy.data.materials.new(name)
material.use_nodes = True
@@ -1347,7 +1349,10 @@ def create_material(
alpha = color_texture.outputs["Alpha"]
material_inputs = nodes.new("ShaderNodeGroup")
is_depth_plane = displacement_image is not None
material_inputs.node_tree = material_node_group(is_depth_plane)
material_inputs.node_tree = material_node_group(
is_depth_plane,
surface_group_name=surface_group_name,
)
material_inputs.name = material_inputs.node_tree.name
material_inputs.label = material_inputs.name
material_inputs.location = (-300.0, 350.0)
@@ -2232,35 +2237,21 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
)
def invoke(self, context, _event):
global _last_depth_plane_model
if not self.properties.is_property_set("snapform_model"):
available_keys = [
item[0] for item in image_depth_model_items(self, context)
]
if available_keys:
self.snapform_model = (
_last_depth_plane_model
if _last_depth_plane_model in available_keys
else available_keys[0]
)
return context.window_manager.invoke_props_dialog(self, width=340)
def draw(self, context):
layout = self.layout
layout.prop(self, "snapform_model")
if self.snapform_model in {"SMALL", "BASE"}:
estimate_box = layout.box()
estimate_box.prop(self, "snapform_use_ray_pose")
draw_process_resolution(estimate_box, self, context.object.data)
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
moge_box = layout.box()
if not model_ready(self.snapform_model):
moge_box.label(text="MoGe-2 model is not installed", icon="ERROR")
moge_box.label(text="Download it in Models")
moge_box.prop(self, "snapform_moge_resolution_level")
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")
layout.separator(type="LINE")
@@ -2369,8 +2360,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
global _last_depth_plane_model
_last_depth_plane_model = self.snapform_model
return {"FINISHED"}

291
src/snapform/image_mesh.py
View File

@@ -3,9 +3,9 @@ from pathlib import Path
import bpy
from .image_depth import (
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME,
create_material,
draw_process_resolution,
image_depth_model_items,
is_image_empty,
plane_bounds,
prepare_image_input,
@@ -16,8 +16,6 @@ from .image_depth import (
)
_last_mesh_model = "BASE"
POINT_ATTRIBUTES = {
"o_depth": ("FLOAT", "point_depth"),
"o_confidence": ("FLOAT", "point_confidence"),
@@ -32,19 +30,28 @@ FACE_ATTRIBUTES = {
"FLOAT",
"face_confidence_percentile",
),
"o_edge_stretch": ("FLOAT", "face_edge_stretch"),
"o_background_face": ("BOOLEAN", "face_background"),
}
OPTIONAL_FACE_ATTRIBUTES = {
"o_mask": ("BOOLEAN", "face_mask"),
}
NODE_GROUP_NAME = "O DA3 Surface"
MOGE_NODE_GROUP_NAME = "O MoGe2 Surface"
NODE_ASSET_NAME = "O_SnapForm.blend"
def load_image_mesh_data(npz_path):
def image_surface_model_items(owner, context):
from .models import surface_model_enum_items
return surface_model_enum_items(owner, context)
def load_image_mesh_data(npz_path, moge2=False):
import numpy as np
try:
with np.load(npz_path, allow_pickle=False) as archive:
required = (
required = [
"vertices",
"faces",
"uv",
@@ -52,13 +59,13 @@ def load_image_mesh_data(npz_path):
"image_size",
"reference_depth",
"point_depth",
"point_confidence",
"point_confidence_percentile",
"point_background",
"face_confidence_percentile",
"face_edge_stretch",
"face_background",
)
]
if moge2:
required.append("face_mask")
else:
required.extend(
("point_confidence", "point_confidence_percentile", "point_background", "face_confidence_percentile", "face_background")
)
missing = [key for key in required if key not in archive.files]
if missing:
raise ValueError(
@@ -81,7 +88,16 @@ def load_image_mesh_data(npz_path):
**POINT_ATTRIBUTES,
**FACE_ATTRIBUTES,
}.items()
if key in archive.files
}
for _name, (_data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items():
if key in archive.files:
attributes[key] = np.asarray(archive[key])
if "vertex_normal" in archive.files:
attributes["vertex_normal"] = np.asarray(
archive["vertex_normal"],
dtype=np.float32,
)
except OSError as error:
raise ValueError(f"Unable to read surface data: {npz_path}") from error
if vertices.ndim != 2 or vertices.shape[1] != 3:
@@ -97,10 +113,13 @@ def load_image_mesh_data(npz_path):
if not reference_depth > 0.0:
raise ValueError("Surface reference depth must be positive")
for _name, (_data_type, key) in POINT_ATTRIBUTES.items():
if len(attributes[key]) != len(vertices):
if key in attributes and len(attributes[key]) != len(vertices):
raise ValueError(f"Point attribute has the wrong length: {key}")
for _name, (_data_type, key) in FACE_ATTRIBUTES.items():
if len(attributes[key]) != len(faces):
if key in attributes and len(attributes[key]) != len(faces):
raise ValueError(f"Face attribute has the wrong length: {key}")
for _name, (_data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items():
if key in attributes and len(attributes[key]) != len(faces):
raise ValueError(f"Face attribute has the wrong length: {key}")
return (
vertices,
@@ -185,7 +204,11 @@ def assign_mesh_attribute(mesh, name, data_type, domain, values):
attribute.data.foreach_set("value", array.reshape(-1))
def assign_snapform_attributes(mesh, attributes):
def assign_snapform_attributes(mesh, attributes, moge2=False):
if moge2:
assign_mesh_attribute(mesh, "o_depth", "FLOAT", "POINT", attributes["point_depth"])
assign_mesh_attribute(mesh, "o_mask", "BOOLEAN", "FACE", attributes["face_mask"])
return
for name, (data_type, key) in POINT_ATTRIBUTES.items():
assign_mesh_attribute(
mesh,
@@ -202,6 +225,36 @@ def assign_snapform_attributes(mesh, attributes):
"FACE",
attributes[key],
)
for name, (data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items():
if key in attributes:
assign_mesh_attribute(
mesh,
name,
data_type,
"FACE",
attributes[key],
)
assign_edge_depth_difference(mesh, attributes["point_depth"])
def assign_edge_depth_difference(mesh, point_depth):
import numpy as np
edge_vertices = np.empty(len(mesh.edges) * 2, dtype=np.int32)
mesh.edges.foreach_get("vertices", edge_vertices)
edge_vertices = edge_vertices.reshape(-1, 2)
depth = np.asarray(point_depth, dtype=np.float32)
first = depth[edge_vertices[:, 0]]
second = depth[edge_vertices[:, 1]]
denominator = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
difference = np.abs(first - second) / denominator
assign_mesh_attribute(
mesh,
"o_edge_stretch",
"FLOAT",
"EDGE",
difference,
)
def remove_surface_background(vertices, faces, uv, attributes):
@@ -228,12 +281,100 @@ def remove_surface_background(vertices, faces, uv, attributes):
)
def prepare_moge_surface_archive(results_path, output_path):
import numpy as np
with np.load(results_path, allow_pickle=False) as archive:
required = {"depth", "mask", "points", "intrinsics"}
missing = required.difference(archive.files)
if missing:
raise ValueError(f"MoGe-2 surface data is missing: {', '.join(sorted(missing))}")
depth = np.asarray(archive["depth"], dtype=np.float32)
mask = np.asarray(archive["mask"], dtype=bool)
points = np.asarray(archive["points"], dtype=np.float32)
intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32)
normal = (
np.asarray(archive["normal"], dtype=np.float32)
if "normal" in archive.files
else None
)
if depth.shape[0] != 1 or points.shape[:3] != (*depth.shape,):
raise ValueError("Convert to Surface requires one matching MoGe-2 point grid")
raw_depth = depth[0]
finite = np.isfinite(raw_depth) & (raw_depth > 0.0)
finite &= np.isfinite(points[0]).all(axis=-1)
height, width = raw_depth.shape
indices = np.arange(height * width, dtype=np.int32).reshape(height, width)
quads = np.stack(
(indices[:-1, :-1], indices[1:, :-1], indices[1:, 1:], indices[:-1, 1:]),
axis=-1,
).reshape(-1, 4)
faces = np.concatenate((quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), axis=0)
faces = faces[finite.reshape(-1)[faces].all(axis=1)]
if not len(faces):
raise ValueError("MoGe-2 output does not contain a valid surface")
flat_points = points[0].reshape(-1, 3)
clean_mask = mask[0] & ~depth_discontinuity_mask(raw_depth, 0.04)
face_mask = clean_mask.reshape(-1)[faces].all(axis=1)
used = np.unique(faces)
index_map = np.full(height * width, -1, dtype=np.int32)
index_map[used] = np.arange(len(used), dtype=np.int32)
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
uv = np.stack(np.meshgrid(u, v), axis=-1).reshape(-1, 2)[used]
camera_intrinsics = intrinsics[0].copy()
principal_x = float(camera_intrinsics[0, 2])
principal_y = float(camera_intrinsics[1, 2])
if 0.0 <= principal_x <= 1.0 and 0.0 <= principal_y <= 1.0:
camera_intrinsics[0] *= width
camera_intrinsics[1] *= height
camera_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = raw_depth.reshape(-1)[used]
masked_depth = raw_depth[finite & mask[0]]
reference_depth = float(np.median(masked_depth if len(masked_depth) else point_depth))
np.savez(
output_path,
vertices=flat_points[used],
faces=index_map[faces],
uv=uv,
intrinsics=camera_intrinsics,
image_size=np.array((width, height), dtype=np.int32),
reference_depth=np.float32(reference_depth),
point_depth=point_depth,
vertex_normal=(
normal[0].reshape(-1, 3)[used]
if normal is not None
else np.empty((0, 3), dtype=np.float32)
),
face_mask=face_mask,
)
return output_path
def depth_discontinuity_mask(depth, threshold):
import numpy as np
edges = np.zeros(depth.shape, dtype=bool)
for axis in (0, 1):
first = np.take(depth, range(depth.shape[axis] - 1), axis=axis)
second = np.take(depth, range(1, depth.shape[axis]), axis=axis)
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
discontinuity = np.abs(first - second) / scale > threshold
lower = [slice(None), slice(None)]
upper = [slice(None), slice(None)]
lower[axis] = slice(0, -1)
upper[axis] = slice(1, None)
edges[tuple(lower)] |= discontinuity
edges[tuple(upper)] |= discontinuity
return edges
def node_asset_path():
return Path(__file__).with_name("assets") / NODE_ASSET_NAME
def load_geometry_node_group():
existing = bpy.data.node_groups.get(NODE_GROUP_NAME)
def load_geometry_node_group(group_name=NODE_GROUP_NAME):
existing = bpy.data.node_groups.get(group_name)
if existing is not None:
return existing
asset_path = node_asset_path()
@@ -243,15 +384,15 @@ def load_geometry_node_group():
if not use_pack:
directory = str(asset_path) + "/NodeTree/"
result = bpy.ops.wm.append(
filepath=directory + NODE_GROUP_NAME,
filepath=directory + group_name,
directory=directory,
filename=NODE_GROUP_NAME,
filename=group_name,
link=False,
do_reuse_local_id=True,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to append the Geometry Nodes asset")
node_group = bpy.data.node_groups.get(NODE_GROUP_NAME)
node_group = bpy.data.node_groups.get(group_name)
if node_group is None:
raise RuntimeError("Appended Geometry Nodes asset is unavailable")
node_group.use_fake_user = True
@@ -260,11 +401,11 @@ def load_geometry_node_group():
source,
target,
):
if NODE_GROUP_NAME not in source.node_groups:
if group_name not in source.node_groups:
raise RuntimeError(
f"Geometry Nodes asset is missing {NODE_GROUP_NAME}"
f"Geometry Nodes asset is missing {group_name}"
)
target.node_groups = [NODE_GROUP_NAME]
target.node_groups = [group_name]
node_group = target.node_groups[0]
if use_pack:
linked_group = node_group
@@ -307,8 +448,13 @@ def material_depth_scale(mesh_object, camera_parameters):
)
def add_geometry_nodes_processing(mesh_object, vertices, faces):
node_group = load_geometry_node_group()
def add_geometry_nodes_processing(
mesh_object,
vertices,
faces,
group_name=NODE_GROUP_NAME,
):
node_group = load_geometry_node_group(group_name)
modifier = mesh_object.modifiers.new(name=node_group.name, type="NODES")
modifier.node_group = node_group
modifier.show_group_selector = False
@@ -360,7 +506,10 @@ def build_image_mesh(
create_camera,
unshaded,
remove_background=False,
moge2=False,
):
import numpy as np
(
camera_vertices,
faces,
@@ -369,7 +518,7 @@ def build_image_mesh(
image_size,
reference_depth,
attributes,
) = load_image_mesh_data(npz_path)
) = load_image_mesh_data(npz_path, moge2=moge2)
if remove_background:
camera_vertices, faces, uv, attributes = remove_surface_background(
camera_vertices,
@@ -388,10 +537,24 @@ def build_image_mesh(
mesh = bpy.data.meshes.new(source_name)
mesh.from_pydata(vertices.tolist(), [], faces.tolist())
assign_vertex_uv(mesh, uv)
assign_snapform_attributes(mesh, attributes)
assign_snapform_attributes(mesh, attributes, moge2=moge2)
mesh.update()
for polygon in mesh.polygons:
polygon.use_smooth = True
if moge2 and len(attributes.get("vertex_normal", ())) == len(vertices):
normals = attributes["vertex_normal"].copy()
normals[:, 0] /= camera_parameters["scale_x"]
normals[:, 1] /= -camera_parameters["scale_y"]
normals[:, 2] /= -camera_parameters["scale_z"]
lengths = np.linalg.norm(normals, axis=1, keepdims=True)
normals /= np.maximum(lengths, 1e-6)
set_custom_normals = getattr(
mesh,
"normals_split_custom_set_from_vertices",
None,
)
if callable(set_custom_normals):
set_custom_normals(normals.tolist())
mesh_object = bpy.data.objects.new(source_name, mesh)
context.collection.objects.link(mesh_object)
mesh_object.matrix_world = source_object.matrix_world.copy()
@@ -401,11 +564,19 @@ def build_image_mesh(
source_object.data,
depth_scale,
unshaded=unshaded,
surface_group_name=(
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME if moge2 else None
),
)
mesh.materials.append(material)
mesh_object["snapform_reference_depth"] = reference_depth
mesh_object["snapform_source_type"] = "IMAGE_MESH"
add_geometry_nodes_processing(mesh_object, vertices, faces)
add_geometry_nodes_processing(
mesh_object,
vertices,
faces,
MOGE_NODE_GROUP_NAME if moge2 else NODE_GROUP_NAME,
)
if create_camera:
create_matching_camera(
context,
@@ -453,6 +624,10 @@ class BuildImageMesh(bpy.types.Operator):
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_moge2: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
@@ -468,6 +643,7 @@ class BuildImageMesh(bpy.types.Operator):
self.snapform_create_camera,
self.snapform_unshaded,
self.snapform_remove_background,
self.snapform_moge2,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
@@ -484,7 +660,8 @@ class ConvertImageToMesh(bpy.types.Operator):
snapform_model: bpy.props.EnumProperty(
name="Model",
description="Model used to estimate image depth",
items=image_depth_model_items,
items=image_surface_model_items,
default=1,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
@@ -492,6 +669,13 @@ class ConvertImageToMesh(bpy.types.Operator):
default=True,
)
snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_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_create_camera: bpy.props.BoolProperty(
name="Create Matching Camera",
description="Create a perspective camera matching the source projection",
@@ -515,7 +699,7 @@ class ConvertImageToMesh(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .models import MODELS
from .models import MODELS, MOGE2_MODELS
from .operators import active_job
from .client import service_busy, model_ready
@@ -524,29 +708,23 @@ class ConvertImageToMesh(bpy.types.Operator):
and video_total_frames(context.object.data) == 1
and active_job() is None
and not service_busy()
and any(model_ready(model_key) for model_key in MODELS)
and any(
model_ready(model_key)
for model_key in (*MOGE2_MODELS, *MODELS)
)
)
def invoke(self, context, _event):
global _last_mesh_model
if not self.properties.is_property_set("snapform_model"):
available_keys = [
item[0] for item in image_depth_model_items(self, context)
]
if available_keys:
self.snapform_model = (
_last_mesh_model
if _last_mesh_model in available_keys
else available_keys[0]
)
return context.window_manager.invoke_props_dialog(self, width=360)
def draw(self, context):
layout = self.layout
layout.prop(self, "snapform_model")
layout.prop(self, "snapform_use_ray_pose")
draw_process_resolution(layout, self, context.object.data)
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")
layout.prop(self, "snapform_unshaded")
@@ -554,17 +732,19 @@ class ConvertImageToMesh(bpy.types.Operator):
layout.prop(self, "snapform_keep_original")
def execute(self, context):
from .operators import IMAGE_MESH_JOB
from .operators import IMAGE_MESH_JOB, MOGE2_MESH_JOB
source_object = context.object
if video_total_frames(source_object.data) != 1:
self.report({"ERROR"}, "Convert to Surface currently supports still images")
return {"CANCELLED"}
try:
process_res = scaled_process_resolution(
source_object.data,
self.snapform_process_res_factor,
)
process_res = 1
if not self.snapform_model.startswith("MOGE2_"):
process_res = scaled_process_resolution(
source_object.data,
self.snapform_process_res_factor,
)
source_path, delete_input = prepare_image_input(
source_object.data,
context.scene,
@@ -573,12 +753,17 @@ class ConvertImageToMesh(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.run_job(
job_type=IMAGE_MESH_JOB,
job_type=(
MOGE2_MESH_JOB
if self.snapform_model.startswith("MOGE2_")
else IMAGE_MESH_JOB
),
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_model_key=self.snapform_model,
snapform_use_ray_pose=self.snapform_use_ray_pose,
snapform_process_res=process_res,
snapform_moge_resolution_level=self.snapform_moge_resolution_level,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_mesh_create_camera=self.snapform_create_camera,
@@ -589,6 +774,4 @@ class ConvertImageToMesh(bpy.types.Operator):
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
global _last_mesh_model
_last_mesh_model = self.snapform_model
return {"FINISHED"}

24
src/snapform/models.py
View File

@@ -139,6 +139,8 @@ MODEL_NUMBERS = {
_ENUM_ITEMS_CACHE = {}
_DEPTH_PLANE_ENUM_ITEMS_CACHE = {}
_SURFACE_ENUM_ITEMS_CACHE = {}
def enum_items(supports_gs=None, show_download_status=False):
@@ -213,9 +215,13 @@ def depth_plane_model_enum_items(_owner, _context):
*MOGE2_MODELS.values(),
*(model for model in MODELS.values() if model.enabled),
]
ready = tuple(model_ready(model.key) for model in models)
cached = _DEPTH_PLANE_ENUM_ITEMS_CACHE.get(ready)
if cached is not None:
return cached
items = []
for index, model in enumerate(models):
label = model.label if model_ready(model.key) else f"{model.label} (Missing)"
label = model.label if ready[index] else f"{model.label} (Missing)"
items.append(
(
model.key,
@@ -225,6 +231,22 @@ def depth_plane_model_enum_items(_owner, _context):
)
)
items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items)))
_DEPTH_PLANE_ENUM_ITEMS_CACHE[ready] = items
return items
def surface_model_enum_items(owner, context):
depth_items = depth_plane_model_enum_items(owner, context)
cache_key = id(depth_items)
cached = _SURFACE_ENUM_ITEMS_CACHE.get(cache_key)
if cached is not None:
return cached
items = [
item
for item in depth_items
if item[0] != "ORIGINAL"
]
_SURFACE_ENUM_ITEMS_CACHE[cache_key] = items
return items

63
src/snapform/operators.py
View File

@@ -20,6 +20,7 @@ from .models import (
RECONSTRUCTION_JOB = "RECONSTRUCT"
IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
MOGE2_DEPTH_JOB = "MOGE2_DEPTH"
MOGE2_MESH_JOB = "MOGE2_MESH"
MOGE2_DEBUG_JOB = "MOGE2_DEBUG"
IMAGE_MESH_JOB = "IMAGE_MESH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
@@ -31,6 +32,7 @@ IMAGE_EMPTY_JOBS = frozenset(
{
*DA3_IMAGE_JOBS,
MOGE2_DEPTH_JOB,
MOGE2_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
@@ -40,6 +42,7 @@ PRODUCTION_AI_JOBS = frozenset(
{
*DA3_IMAGE_JOBS,
MOGE2_DEPTH_JOB,
MOGE2_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}
@@ -163,6 +166,7 @@ class RunJob(bpy.types.Operator):
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(MOGE2_DEPTH_JOB, "Convert Image with MoGe-2", ""),
(MOGE2_MESH_JOB, "Convert Image to Surface with MoGe-2", ""),
(MOGE2_DEBUG_JOB, "Test MoGe-2 API", ""),
(IMAGE_MESH_JOB, "Convert Image to Surface", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
@@ -250,7 +254,7 @@ class RunJob(bpy.types.Operator):
export_format = "mini_npz"
elif self.job_type == IMAGE_MESH_JOB:
export_format = "mesh_data"
elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}:
elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}:
export_format = "moge_npz"
job = ActiveJob(
@@ -425,8 +429,12 @@ class RunJob(bpy.types.Operator):
),
paths,
)
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}:
model_key = self.snapform_model_key or MOGE2_MODEL.key
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}:
model_key = (
settings.moge_model
if self.job_type == MOGE2_DEBUG_JOB
else self.snapform_model_key or MOGE2_MODEL.key
)
model = get_downloadable_model(model_key)
if not client.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
@@ -440,17 +448,23 @@ class RunJob(bpy.types.Operator):
input_path=input_path,
resolution_level=(
self.snapform_moge_resolution_level
if self.job_type == MOGE2_DEPTH_JOB
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB}
else settings.moge_resolution_level
),
include_points=self.job_type == MOGE2_DEBUG_JOB,
include_points=self.job_type in {MOGE2_MESH_JOB, MOGE2_DEBUG_JOB},
remove_background=(
self.job_type == MOGE2_DEPTH_JOB
self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB}
and remove_background
),
start_frame=start_frame,
device=self._task_device(settings),
model_key=model_key,
debug_export_format=(
settings.moge_export_format
if self.job_type == MOGE2_DEBUG_JOB
else ""
),
surface_mesh=self.job_type == MOGE2_MESH_JOB,
),
paths,
)
@@ -531,7 +545,7 @@ class RunJob(bpy.types.Operator):
return f"Starting {settings.export_format} export"
if self.job_type == IMAGE_DEPTH_JOB:
return "Starting depth estimation"
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}:
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}:
return "Starting MoGe-2 inference"
if self.job_type == IMAGE_MESH_JOB:
return "Starting perspective surface generation"
@@ -728,6 +742,21 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "MoGe-2 depth plane is ready"
if job.kind == MOGE2_MESH_JOB:
try:
settings.status = "Building MoGe-2 surface"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_moge2_image_mesh(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "MoGe-2 surface import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "MoGe-2 surface is ready"
if job.kind == IMAGE_MESH_JOB:
try:
settings.status = "Building perspective surface"
@@ -801,6 +830,26 @@ class RunJob(bpy.types.Operator):
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the surface Undo step")
def _build_moge2_image_mesh(self, job):
result_directory = job.result_directory
mesh_path = result_directory / "mesh.npz"
if not mesh_path.is_file():
raise RuntimeError("MoGe-2 did not export the official surface mesh")
result = bpy.ops.snapform.build_image_mesh(
snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(mesh_path),
snapform_delete_source=job.delete_source,
snapform_create_camera=job.mesh_create_camera,
snapform_unshaded=job.unshaded,
snapform_remove_background=False,
snapform_moge2=True,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the MoGe-2 perspective surface")
undo_result = bpy.ops.ed.undo_push(message="Convert to Surface")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the surface Undo step")
def _build_depth_plane(self, context, job):
result_directory = job.result_directory
npz_path = result_directory / "results.npz"

3
src/snapform/panel.py
View File

@@ -252,9 +252,10 @@ class Moge2ApiDebugPanel(bpy.types.Panel):
def draw(self, context):
layout = self.layout
settings = context.scene.snapform_settings
layout.prop(settings, "moge_export_format")
layout.prop(settings, "input_path")
layout.prop(settings, "moge_resolution_level")
layout.prop(settings, "moge_model")
layout.prop(settings, "moge_resolution_level")
run_row = layout.row()
run_row.enabled = (
active_job() is None

10
src/snapform/properties.py
View File

@@ -78,6 +78,16 @@ class SnapFormSettings(bpy.types.PropertyGroup):
items=moge2_model_enum_items,
default=1,
)
moge_export_format: bpy.props.EnumProperty(
name="Output Format",
description="Additional MoGe-2 debug output generated with results.npz",
items=(
("maps", "Maps", "Export depth, normal, and mask previews"),
("glb", "Textured Surface (GLB)", "Export the official-style textured mesh"),
("ply", "Point Cloud (PLY)", "Export the official-style colored point cloud"),
),
default="maps",
)
conf_thresh_percentile: bpy.props.FloatProperty(
name="Confidence Threshold",
description="Remove the lowest-confidence percentile of points or surface faces",

141
src/snapform/worker.py
View File

@@ -1514,6 +1514,94 @@ def _save_moge_preview(output_directory, index, depth, normal, mask):
)
def export_moge_debug_mesh(
output_directory,
points,
depth,
mask,
normal,
image,
export_format,
):
import numpy as np
import utils3d
from moge.utils.io import save_glb, save_ply
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(
points,
image,
uv,
normal,
mask=clean_mask,
tri=True,
)
faces, vertices, colors, vertex_uvs, vertex_normals = result
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,
)
def export_moge_surface_archive(
output_path,
points,
depth,
mask,
normal,
intrinsics,
):
import numpy as np
import utils3d
height, width = depth.shape
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
result = utils3d.np.build_mesh_from_map(
points,
utils3d.np.uv_map(height, width),
normal,
mask=clean_mask,
tri=True,
)
faces, vertices, vertex_uvs, vertex_normals = result
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
pixel_intrinsics = intrinsics.copy()
pixel_intrinsics[0] *= width
pixel_intrinsics[1] *= height
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = vertices[:, 2].astype(np.float32)
np.savez(
output_path,
vertices=vertices.astype(np.float32),
faces=faces.astype(np.int32),
uv=vertex_uvs.astype(np.float32),
intrinsics=pixel_intrinsics.astype(np.float32),
image_size=np.array((width, height), dtype=np.int32),
reference_depth=np.float32(np.median(point_depth)),
point_depth=point_depth,
vertex_normal=vertex_normals.astype(np.float32),
face_mask=np.ones(len(faces), dtype=bool),
)
def run_moge(args, model_cache=None, cancel_check=None):
import numpy as np
from PIL import Image
@@ -1550,6 +1638,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
masks = []
points = []
intrinsics = []
source_images = []
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
write_status(
@@ -1560,6 +1649,7 @@ 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,
@@ -1579,13 +1669,15 @@ def run_moge(args, model_cache=None, cancel_check=None):
intrinsics.append(intrinsic)
if bool(getattr(args, "include_points", False)):
points.append(point)
_save_moge_preview(
output_directory,
None if len(images) == 1 else index,
depth,
normal,
mask,
)
debug_format = getattr(args, "debug_export_format", "")
if not debug_format or debug_format == "maps":
_save_moge_preview(
output_directory,
None if len(images) == 1 else index,
depth,
normal,
mask,
)
shapes = {array.shape for array in depths}
if len(shapes) != 1:
@@ -1598,7 +1690,18 @@ def run_moge(args, model_cache=None, cancel_check=None):
}
if points:
result["points"] = np.stack(points).astype(np.float32)
np.savez(output_directory / "results.npz", **result)
debug_format = getattr(args, "debug_export_format", "")
if debug_format in {"glb", "ply"} and len(images) == 1:
write_status(args.status, 0.84, f"Exporting MoGe-2 {debug_format.upper()}")
export_moge_debug_mesh(
output_directory,
points[0],
depths[0],
masks[0],
normals[0],
source_images[0],
debug_format,
)
remove_background_model_dir = getattr(
args, "remove_background_model_dir", ""
@@ -1614,7 +1717,14 @@ def run_moge(args, model_cache=None, cancel_check=None):
model_cache=model_cache,
cancel_check=cancel_check,
)
foreground.save(output_directory / "foreground.png")
foreground_path = output_directory / "foreground.png"
foreground.save(foreground_path)
if points:
foreground_mask = foreground_mesh_mask(
foreground_path,
result["depth"][0].shape,
)
result["mask"][0] &= foreground_mask
else:
infer_foreground_frames(
images,
@@ -1628,6 +1738,19 @@ def run_moge(args, model_cache=None, cancel_check=None):
model_cache=model_cache,
cancel_check=cancel_check,
)
if bool(getattr(args, "surface_mesh", False)):
if len(images) != 1 or not points:
raise RuntimeError("MoGe-2 Surface requires one point map")
write_status(args.status, 0.94, "Building official MoGe-2 surface")
export_moge_surface_archive(
output_directory / "mesh.npz",
points[0],
result["depth"][0],
result["mask"][0],
result["normal"][0],
result["intrinsics"][0],
)
np.savez(output_directory / "results.npz", **result)
copy_frames(images, output_directory)
write_status(
args.status,