from pathlib import Path import bpy from .image_depth import ( MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, create_material, draw_process_resolution, is_image_empty, normal_data_name, plane_bounds, prepare_image_input, process_resolution_scale_property, remove_temporary_input, scaled_process_resolution, video_total_frames, ) POINT_ATTRIBUTES = { "o_depth": ("FLOAT", "point_depth"), } FACE_ATTRIBUTES = { "o_face_confidence": ( "FLOAT", "face_confidence_percentile", ), } OPTIONAL_FACE_ATTRIBUTES = { "o_mask": ("BOOLEAN", "face_mask"), "o_depth_discontinuity": ("FLOAT", "face_depth_discontinuity"), } NODE_GROUP_NAME = "O DA3 Surface" MOGE_NODE_GROUP_NAME = "O MoGe2 Surface" NODE_ASSET_NAME = "O_SnapForm.blend" 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 = [ "vertices", "faces", "uv", "intrinsics", "image_size", "reference_depth", "point_depth", ] if moge2: required.extend(("face_mask", "face_depth_discontinuity", "normal_map")) else: required.extend( ("face_confidence_percentile", "face_depth_discontinuity") ) missing = [key for key in required if key not in archive.files] if missing: raise ValueError( f"Surface data is missing: {', '.join(missing)}" ) vertices = np.asarray(archive["vertices"], dtype=np.float32) faces = np.asarray(archive["faces"], dtype=np.int32) uv = np.asarray(archive["uv"], dtype=np.float32) intrinsics = np.asarray( archive["intrinsics"], dtype=np.float32, ) image_size = tuple( int(value) for value in archive["image_size"].tolist() ) reference_depth = float(archive["reference_depth"]) attributes = { key: np.asarray(archive[key]) for _name, (_data_type, key) in { **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 "normal_map" in archive.files: attributes["normal_map"] = np.asarray( archive["normal_map"], 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: raise ValueError(f"Unsupported vertices shape: {vertices.shape}") if faces.ndim != 2 or faces.shape[1] != 3: raise ValueError(f"Unsupported faces shape: {faces.shape}") if uv.shape != (len(vertices), 2): raise ValueError("Surface UV and vertex counts do not match") if intrinsics.shape != (3, 3): raise ValueError(f"Unsupported intrinsics shape: {intrinsics.shape}") if len(image_size) != 2 or min(image_size) <= 0: raise ValueError(f"Unsupported image size: {image_size}") if not reference_depth > 0.0: raise ValueError("Surface reference depth must be positive") if moge2 and attributes["normal_map"].shape != ( image_size[1], image_size[0], 3, ): raise ValueError("Surface normal map and image sizes do not match") if moge2 and not np.isfinite(attributes["normal_map"]).all(): raise ValueError("Surface normal map contains invalid values") for _name, (_data_type, key) in POINT_ATTRIBUTES.items(): 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 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, faces, uv, intrinsics, image_size, reference_depth, attributes, ) def empty_anchored_vertices( source_object, camera_vertices, intrinsics, image_size, reference_depth, ): import numpy as np width, height = image_size x_min, x_max, y_min, y_max = plane_bounds(source_object) plane_width = x_max - x_min plane_height = y_max - y_min pixel_width = max(width - 1, 1) pixel_height = max(height - 1, 1) fx = float(intrinsics[0, 0]) fy = float(intrinsics[1, 1]) cx = float(intrinsics[0, 2]) cy = float(intrinsics[1, 2]) if fx <= 0.0 or fy <= 0.0: raise ValueError("Camera intrinsics contain an invalid focal length") scale_x = plane_width * fx / (pixel_width * reference_depth) scale_y = plane_height * fy / (pixel_height * reference_depth) scale_z = float(np.sqrt(scale_x * scale_y)) camera_x = x_min + plane_width * cx / pixel_width camera_y = y_max - plane_height * cy / pixel_height vertices = np.empty_like(camera_vertices, dtype=np.float32) vertices[:, 0] = camera_x + camera_vertices[:, 0] * scale_x vertices[:, 1] = camera_y - camera_vertices[:, 1] * scale_y vertices[:, 2] = ( reference_depth - camera_vertices[:, 2] ) * scale_z camera_location = ( camera_x, camera_y, reference_depth * scale_z, ) camera_parameters = { "scale_x": scale_x, "scale_y": scale_y, "scale_z": scale_z, "location": camera_location, } return vertices, camera_parameters def assign_vertex_uv(mesh, vertex_uv): import numpy as np uv_layer = mesh.uv_layers.new(name="UVMap") loop_vertices = np.empty(len(mesh.loops), dtype=np.int32) mesh.loops.foreach_get("vertex_index", loop_vertices) loop_uv = vertex_uv[loop_vertices] uv_layer.data.foreach_set("uv", loop_uv.reshape(-1)) def assign_mesh_attribute(mesh, name, data_type, domain, values): import numpy as np attribute = mesh.attributes.new( name=name, type=data_type, domain=domain, ) array = np.asarray(values) if data_type == "BOOLEAN": array = array.astype(bool) else: array = array.astype(np.float32) attribute.data.foreach_set("value", array.reshape(-1)) 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"]) assign_mesh_attribute( mesh, "o_depth_discontinuity", "FLOAT", "FACE", attributes["face_depth_discontinuity"], ) return for name, (data_type, key) in POINT_ATTRIBUTES.items(): assign_mesh_attribute( mesh, name, data_type, "POINT", attributes[key], ) for name, (data_type, key) in FACE_ATTRIBUTES.items(): assign_mesh_attribute( mesh, name, data_type, "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_mesh_attribute( mesh, "o_depth_discontinuity", "FLOAT", "FACE", attributes["face_depth_discontinuity"], ) 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", "normal", "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 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) face_mask = mask[0].reshape(-1)[faces].all(axis=1) depth_edge_ratio = depth_discontinuity_ratio(raw_depth) face_depth_discontinuity = depth_edge_ratio.reshape(-1)[faces].max(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, normal_map=normal[0], face_mask=face_mask, face_depth_discontinuity=face_depth_discontinuity, ) return output_path def depth_discontinuity_ratio(depth): import numpy as np padded = np.pad(depth, 1, mode="edge") windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3)) difference = np.nanmax(windows, axis=(-2, -1)) difference -= np.nanmin(windows, axis=(-2, -1)) return difference / np.maximum(depth, 1e-6) def node_asset_path(): return Path(__file__).with_name("assets") / NODE_ASSET_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() if not asset_path.is_file(): raise RuntimeError(f"Missing Geometry Nodes asset: {asset_path}") use_pack = bpy.app.version >= (5, 0, 0) if not use_pack: directory = str(asset_path) + "/NodeTree/" result = bpy.ops.wm.append( filepath=directory + group_name, directory=directory, 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(group_name) if node_group is None: raise RuntimeError("Appended Geometry Nodes asset is unavailable") node_group.use_fake_user = True return node_group with bpy.data.libraries.load(str(asset_path), link=use_pack) as ( source, target, ): if group_name not in source.node_groups: raise RuntimeError( f"Geometry Nodes asset is missing {group_name}" ) target.node_groups = [group_name] node_group = target.node_groups[0] if use_pack: linked_group = node_group node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group) if linked_group != node_group and linked_group.users == 0: bpy.data.node_groups.remove(linked_group) node_group.use_fake_user = True return node_group def modifier_input_identifier(node_group, name): for item in node_group.interface.items_tree: if item.item_type == "SOCKET" and item.in_out == "INPUT": if item.name == name: return item.identifier raise RuntimeError(f"Geometry Nodes input is missing: {name}") def default_remesh_resolution(vertices, faces): import numpy as np if not len(faces): return 0.01 triangles = np.asarray(vertices, dtype=np.float32)[faces] edge_lengths = np.linalg.norm( triangles - np.roll(triangles, 1, axis=1), axis=2, ) valid = edge_lengths[np.isfinite(edge_lengths) & (edge_lengths > 0.0)] if not len(valid): return 0.01 # Twice the 90% reduction estimate gives the iterative remesh more room # to settle into a smooth, stable default topology. return float(np.median(valid) * np.sqrt(10.0) * 2.0) def material_depth_scale(mesh_object, camera_parameters): return abs(float(mesh_object.scale.z)) * abs( float(camera_parameters["scale_z"]) ) 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 identifier = modifier_input_identifier(node_group, "Resolution") modifier[identifier] = default_remesh_resolution(vertices, faces) def create_matching_camera( context, source_object, mesh_object, name, intrinsics, image_size, camera_parameters, ): from mathutils import Matrix width, height = image_size scale_x = camera_parameters["scale_x"] scale_z = camera_parameters["scale_z"] effective_fx = float(intrinsics[0, 0]) * scale_z / scale_x camera_data = bpy.data.cameras.new(f"{name}.camera") camera_data.sensor_fit = "HORIZONTAL" camera_data.sensor_width = 36.0 camera_data.lens = effective_fx * camera_data.sensor_width / max( width - 1, 1, ) cx = float(intrinsics[0, 2]) cy = float(intrinsics[1, 2]) camera_data.shift_x = (cx - (width - 1) * 0.5) / max(width - 1, 1) camera_data.shift_y = ((height - 1) * 0.5 - cy) / max(width - 1, 1) camera_object = bpy.data.objects.new(f"{name}.camera", camera_data) context.collection.objects.link(camera_object) local_matrix = Matrix.Translation(camera_parameters["location"]) camera_object.matrix_world = source_object.matrix_world @ local_matrix camera_world = camera_object.matrix_world.copy() camera_object.parent = mesh_object camera_object.matrix_world = camera_world return camera_object def build_image_mesh( context, source_object, npz_path, delete_source, create_camera, unshaded, moge2=False, ): import numpy as np ( camera_vertices, faces, uv, intrinsics, image_size, reference_depth, attributes, ) = load_image_mesh_data(npz_path, moge2=moge2) vertices, camera_parameters = empty_anchored_vertices( source_object, camera_vertices, intrinsics, image_size, reference_depth, ) source_name = source_object.name mesh = bpy.data.meshes.new(source_name) mesh.from_pydata(vertices.tolist(), [], faces.tolist()) assign_vertex_uv(mesh, uv) assign_snapform_attributes(mesh, attributes, moge2=moge2) mesh.update() for polygon in mesh.polygons: polygon.use_smooth = True mesh_object = bpy.data.objects.new(source_name, mesh) context.collection.objects.link(mesh_object) mesh_object.matrix_world = source_object.matrix_world.copy() depth_scale = material_depth_scale(mesh_object, camera_parameters) normal_image = None if moge2 and attributes.get("normal_map", np.empty(0)).size: from .image_depth import create_moge_normal_image normal_map = attributes["normal_map"].copy() normal_map[..., 0] /= camera_parameters["scale_x"] normal_map[..., 1] /= -camera_parameters["scale_y"] normal_map[..., 2] /= -camera_parameters["scale_z"] lengths = np.linalg.norm(normal_map, axis=-1, keepdims=True) normal_map /= np.maximum(lengths, 1e-6) normal_image = create_moge_normal_image( normal_data_name(source_object), normal_map, convert_coordinates=False, ) material = create_material( source_name, source_object.data, depth_scale, unshaded=unshaded, surface_group_name=( MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME if moge2 else None ), normal_image=normal_image, ) 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, MOGE_NODE_GROUP_NAME if moge2 else NODE_GROUP_NAME, ) if create_camera: create_matching_camera( context, source_object, mesh_object, source_name, intrinsics, image_size, camera_parameters, ) for selected in context.selected_objects: selected.select_set(False) mesh_object.select_set(True) context.view_layer.objects.active = mesh_object if delete_source: bpy.data.objects.remove(source_object, do_unlink=True) mesh_object.name = source_name return mesh_object class BuildImageMesh(bpy.types.Operator): bl_idname = "snapform.build_image_mesh" bl_label = "Build Image Surface" 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_delete_source: bpy.props.BoolProperty( default=True, options={"HIDDEN", "SKIP_SAVE"}, ) snapform_create_camera: bpy.props.BoolProperty( default=False, options={"HIDDEN", "SKIP_SAVE"}, ) snapform_unshaded: bpy.props.BoolProperty( 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) if not is_image_empty(source_object): self.report({"ERROR"}, "The source Image Empty no longer exists") return {"CANCELLED"} try: build_image_mesh( context, source_object, self.snapform_npz_path, self.snapform_delete_source, self.snapform_create_camera, self.snapform_unshaded, self.snapform_moge2, ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} return {"FINISHED"} class ConvertImageToMesh(bpy.types.Operator): bl_idname = "snapform.convert_image_to_mesh" bl_label = "Convert to Surface" bl_description = "Convert an Image Empty into a perspective textured surface" bl_options = {"REGISTER"} snapform_model: bpy.props.EnumProperty( name="Depth From", description="Model used to estimate image depth", items=image_surface_model_items, default=1, ) snapform_use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", description="Use ray-based camera pose estimation", 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", default=False, ) snapform_unshaded: bpy.props.BoolProperty( name="Unshaded", description="Display the source color through an Emission shader", default=False, ) snapform_remove_background: bpy.props.BoolProperty( name="Remove Background", description="Remove the background during conversion with BEN2", default=False, ) snapform_keep_original: bpy.props.BoolProperty( name="Keep Original", description="Keep the source Image Empty after conversion", default=False, ) @classmethod def poll(cls, context): from .models import MODELS, MOGE2_MODELS from .operators import active_job from .client import service_busy, model_ready return ( is_image_empty(context.object) 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 (*MOGE2_MODELS, *MODELS) ) ) def invoke(self, context, _event): return context.window_manager.invoke_props_dialog(self, width=360) def draw(self, context): layout = self.layout layout.prop(self, "snapform_model") 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") layout.separator(type="LINE") layout.prop(self, "snapform_keep_original") def execute(self, context): 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 = 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, ) except (RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} result = bpy.ops.snapform.run_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, snapform_unshaded=self.snapform_unshaded, snapform_remove_background=self.snapform_remove_background, ) if "CANCELLED" in result: if delete_input: remove_temporary_input(source_path) return {"CANCELLED"} return {"FINISHED"}