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

Add geometry nodes asset and validation script

Browse Source 
- Implemented a new geometry nodes asset in `build_geometry_nodes_asset.py` to handle mesh refinement and topology fixing.
- Created a validation script in `validate_geometry_nodes_asset.py` to ensure the integrity and functionality of the geometry node group.
- Added input handling for various parameters including resolution, iterations, and smoothing options.
- Included checks for node group structure and attributes to validate the geometry processing pipeline.
This commit is contained in:
马南
2026-08-11 13:23:13 +08:00
parent fc27c23fd1
commit 07d836d9ad
15 changed files with 2408 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/snapform/__init__.py
View File

@@ -7,8 +7,10 @@ from .image_depth import (
BuildOriginalDepthPlane,
ConvertImageToDepthPlane,
RemoveImageBackground,
SnapFormImageMenu,
draw_image_context_menu,
)
from .image_mesh import BuildImageMesh, ConvertImageToMesh
from .operators import (
CancelJob,
ControlHost,
@@ -34,11 +36,14 @@ CLASSES = (
RunJob,
CancelJob,
ControlHost,
SnapFormImageMenu,
ApplyRemovedBackground,
RemoveImageBackground,
BuildDepthPlane,
BuildOriginalDepthPlane,
ConvertImageToDepthPlane,
BuildImageMesh,
ConvertImageToMesh,
HostPanel,
ModelsPanel,
DebugPanel,

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

Binary file not shown.

45
src/snapform/image_depth.py
View File

@@ -1783,22 +1783,37 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return {"FINISHED"}
class SnapFormImageMenu(bpy.types.Menu):
bl_idname = "SNAPFORM_MT_image_context"
bl_label = "SnapForm"
def draw(self, _context):
layout = self.layout
layout.operator_context = "INVOKE_DEFAULT"
layout.operator(
"snapform.convert_image_to_depth_plane",
text="Convert to Depth Plane",
icon="MOD_DISPLACE",
)
layout.operator(
"snapform.convert_image_to_mesh",
text="Convert to Mesh",
icon="MESH_GRID",
)
layout.separator()
layout.operator(
"snapform.remove_image_background",
text="Remove Background",
icon="IMAGE_ALPHA",
)
def draw_image_context_menu(self, context):
active_object = context.object
if not is_image_empty(active_object):
if not is_image_empty(context.object):
return
self.layout.separator()
row = self.layout.row()
row.operator_context = "INVOKE_DEFAULT"
row.operator(
"snapform.convert_image_to_depth_plane",
text="Convert to Depth Plane",
icon="MOD_DISPLACE",
)
row = self.layout.row()
row.operator_context = "INVOKE_DEFAULT"
row.operator(
"snapform.remove_image_background",
text="Remove Background",
icon="IMAGE_ALPHA",
self.layout.menu(
SnapFormImageMenu.bl_idname,
text="SnapForm",
icon="PLUGIN",
)

564
src/snapform/image_mesh.py Normal file
View File

@@ -0,0 +1,564 @@
from pathlib import Path
import bpy
from .image_depth import (
create_material,
image_depth_model_items,
is_image_empty,
plane_bounds,
prepare_image_input,
remove_temporary_input,
scaled_process_resolution,
video_total_frames,
)
_last_mesh_model = "BASE"
POINT_ATTRIBUTES = {
"o_depth": ("FLOAT", "point_depth"),
"o_confidence": ("FLOAT", "point_confidence"),
"o_confidence_percentile": (
"FLOAT",
"point_confidence_percentile",
),
"o_background": ("BOOLEAN", "point_background"),
}
FACE_ATTRIBUTES = {
"o_face_confidence": (
"FLOAT",
"face_confidence_percentile",
),
"o_edge_stretch": ("FLOAT", "face_edge_stretch"),
"o_background_face": ("BOOLEAN", "face_background"),
}
NODE_GROUP_NAME = "O Refine DA3 Mesh"
NODE_ASSET_NAME = "O_SnapForm.blend"
def load_image_mesh_data(npz_path):
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",
"point_confidence",
"point_confidence_percentile",
"point_background",
"face_confidence_percentile",
"face_edge_stretch",
"face_background",
)
missing = [key for key in required if key not in archive.files]
if missing:
raise ValueError(
f"Mesh 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()
}
except OSError as error:
raise ValueError(f"Unable to read mesh 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("Mesh 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("Mesh reference depth must be positive")
for _name, (_data_type, key) in POINT_ATTRIBUTES.items():
if 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):
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):
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],
)
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)
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 + NODE_GROUP_NAME,
directory=directory,
filename=NODE_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)
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 NODE_GROUP_NAME not in source.node_groups:
raise RuntimeError(
f"Geometry Nodes asset is missing {NODE_GROUP_NAME}"
)
target.node_groups = [NODE_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):
node_group = load_geometry_node_group()
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,
):
(
camera_vertices,
faces,
uv,
intrinsics,
image_size,
reference_depth,
attributes,
) = load_image_mesh_data(npz_path)
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)
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)
material = create_material(
source_name,
source_object.data,
depth_scale,
unshaded=unshaded,
)
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)
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 Mesh"
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"},
)
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,
)
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 Mesh"
bl_description = "Convert an Image Empty into a perspective textured mesh"
bl_options = {"REGISTER"}
snapform_model: bpy.props.EnumProperty(
name="Model",
description="Model used to estimate image depth",
items=image_depth_model_items,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Use ray-based camera pose estimation",
default=True,
)
snapform_process_res_factor: bpy.props.FloatProperty(
name="Resolution Scale",
description="Scale applied to the source image long-edge resolution",
default=1.0,
min=0.1,
max=1.0,
step=10,
precision=2,
)
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_delete_original: bpy.props.BoolProperty(
name="Delete Original",
description="Delete the source Image Empty after conversion",
default=True,
)
@classmethod
def poll(cls, context):
from .models import BEN2_MODEL, MODELS
from .operators import active_job
from .runtime import host_busy, model_ready
return (
is_image_empty(context.object)
and video_total_frames(context.object.data) == 1
and active_job() is None
and not host_busy()
and any(model_ready(model_key) for model_key in MODELS)
and model_ready(BEN2_MODEL.key)
)
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")
process_res = scaled_process_resolution(
context.object.data,
self.snapform_process_res_factor,
)
layout.prop(
self,
"snapform_process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
layout.prop(self, "snapform_create_camera")
layout.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_delete_original")
def execute(self, context):
from .operators import IMAGE_MESH_JOB
source_object = context.object
if video_total_frames(source_object.data) != 1:
self.report({"ERROR"}, "Convert to Mesh currently supports still images")
return {"CANCELLED"}
try:
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=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_delete_source=self.snapform_delete_original,
snapform_delete_input=delete_input,
snapform_mesh_create_camera=self.snapform_create_camera,
snapform_unshaded=self.snapform_unshaded,
)
if "CANCELLED" in result:
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
global _last_mesh_model
_last_mesh_model = self.snapform_model
return {"FINISHED"}

72
src/snapform/operators.py
View File

@@ -17,6 +17,7 @@ from .models import (
RECONSTRUCTION_JOB = "RECONSTRUCT"
IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
IMAGE_MESH_JOB = "IMAGE_MESH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND"
@@ -78,6 +79,7 @@ class ActiveJob:
unshaded: bool = False
remove_background: bool = False
background_depth_distance: float = 0.0
mesh_create_camera: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
uses_host: bool = False
@@ -139,6 +141,7 @@ class RunJob(bpy.types.Operator):
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(IMAGE_MESH_JOB, "Convert Image to Mesh", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
@@ -203,7 +206,14 @@ class RunJob(bpy.types.Operator):
min=0.0,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_mesh_create_camera: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_remove_background: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
global _active_job
@@ -227,6 +237,8 @@ class RunJob(bpy.types.Operator):
export_format = settings.export_format
elif self.job_type == IMAGE_DEPTH_JOB:
export_format = "mini_npz"
elif self.job_type == IMAGE_MESH_JOB:
export_format = "mesh_data"
job = ActiveJob(
kind=self.job_type,
@@ -243,12 +255,21 @@ class RunJob(bpy.types.Operator):
delete_source=self.snapform_delete_source,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
remove_background=self.snapform_background_depth_distance > 0.0,
remove_background=(
self.snapform_background_depth_distance > 0.0
or bool(getattr(self, "snapform_remove_background", False))
),
background_depth_distance=self.snapform_background_depth_distance,
mesh_create_camera=getattr(
self,
"snapform_mesh_create_camera",
False,
),
temporary_input_path=(
Path(self.snapform_input_path)
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
@@ -317,6 +338,7 @@ class RunJob(bpy.types.Operator):
getattr(self, "snapform_delete_input", False)
and self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
@@ -332,6 +354,8 @@ class RunJob(bpy.types.Operator):
getattr(self, "snapform_background_depth_distance", 0.0)
)
remove_background = background_depth_distance > 0.0
if self.job_type == IMAGE_MESH_JOB:
remove_background = True
if self.job_type == "INSTALL_RUNTIME":
self._require_online_access()
return runtime.install_runtime_command(paths["status"]), paths
@@ -349,6 +373,7 @@ class RunJob(bpy.types.Operator):
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
@@ -410,15 +435,18 @@ class RunJob(bpy.types.Operator):
if self.job_type not in {
RECONSTRUCTION_JOB,
IMAGE_DEPTH_JOB,
IMAGE_MESH_JOB,
}:
raise RuntimeError(f"Unsupported task type: {self.job_type}")
export_format = settings.export_format
if self.job_type == IMAGE_DEPTH_JOB:
export_format = "mini_npz"
elif self.job_type == IMAGE_MESH_JOB:
export_format = "mesh_data"
use_ray_pose = None
process_res = None
model_key = selected_model_key(settings, export_format)
if self.job_type == IMAGE_DEPTH_JOB:
if self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}:
model_key = self.snapform_model_key
use_ray_pose = self.snapform_use_ray_pose
process_res = self.snapform_process_res
@@ -426,7 +454,7 @@ class RunJob(bpy.types.Operator):
if not runtime.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
if (
self.job_type == IMAGE_DEPTH_JOB
self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}
and remove_background
and not runtime.model_ready(BEN2_MODEL.key)
):
@@ -446,13 +474,13 @@ class RunJob(bpy.types.Operator):
process_res=process_res,
process_res_method=(
IMAGE_DEPTH_PROCESS_RES_METHOD
if self.job_type == IMAGE_DEPTH_JOB
if self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}
else None
),
use_ray_pose=use_ray_pose,
start_frame=start_frame,
remove_background=(
self.job_type == IMAGE_DEPTH_JOB
self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}
and remove_background
),
background_depth_distance=(
@@ -469,6 +497,7 @@ class RunJob(bpy.types.Operator):
"""Debug panels use the scene device; production tasks use Preferences."""
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}:
@@ -490,6 +519,8 @@ 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 == IMAGE_MESH_JOB:
return "Starting perspective mesh generation"
if self.job_type == BACKGROUND_REMOVAL_JOB:
return "Starting BEN2 background removal"
if self.job_type == IMAGE_BACKGROUND_JOB:
@@ -663,6 +694,21 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Depth plane is ready"
if job.kind == IMAGE_MESH_JOB:
try:
settings.status = "Building perspective mesh"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_image_mesh(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "Mesh import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Perspective mesh is ready"
if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}:
try:
settings.status = "Building depth plane"
@@ -706,6 +752,20 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"FINISHED"}
def _build_image_mesh(self, job):
result = bpy.ops.snapform.build_image_mesh(
snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(job.result_directory / "mesh.npz"),
snapform_delete_source=job.delete_source,
snapform_create_camera=job.mesh_create_camera,
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the perspective mesh")
undo_result = bpy.ops.ed.undo_push(message="Convert to Mesh")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the mesh Undo step")
def _build_depth_plane(self, context, job):
result_directory = job.result_directory
npz_path = result_directory / "results.npz"

4
src/snapform/runtime.py
View File

@@ -37,6 +37,7 @@ JOB_FOLDER_NAMES = {
"DOWNLOAD_MODEL": "download-model",
"RECONSTRUCT": "export-reconstruction",
"IMAGE_DEPTH": "generate-depth-plane",
"IMAGE_MESH": "generate-image-mesh",
"REMOVE_BACKGROUND": "remove-background",
"IMAGE_BACKGROUND": "remove-image-background",
"ORIGINAL_BACKGROUND": "generate-original-depth-plane",
@@ -812,6 +813,7 @@ def reconstruct_request(
use_ray_pose=None,
remove_background=False,
background_depth_distance=0.0,
mesh_filter_edges=None,
start_frame=1,
device=None,
):
@@ -852,6 +854,8 @@ def reconstruct_request(
"background_depth_distance": float(background_depth_distance),
"mesh_filter_edges": bool(
getattr(settings, "mesh_filter_edges", True)
if mesh_filter_edges is None
else mesh_filter_edges
),
"mesh_max_edge_length": float(
getattr(settings, "mesh_max_edge_length", 0.1)

468
src/snapform/worker.py
View File

@@ -24,7 +24,14 @@ VIDEO_SUFFIXES = {
".mpeg",
}
SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz", "mesh")
EXPORT_FORMATS = (
"glb",
"gs_ply",
"depth_vis",
"mini_npz",
"mesh",
"mesh_data",
)
HF_ENDPOINT = "https://huggingface.co"
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
MODELSCOPE_REVISION = "master"
@@ -368,41 +375,357 @@ def unproject_depth_frame(depth, intrinsics, extrinsics):
return world_points[:, :3].reshape(height, width, 3)
def unproject_camera_depth(depth, intrinsics):
import numpy as np
height, width = depth.shape
x, y = np.meshgrid(
np.arange(width, dtype=np.float32),
np.arange(height, dtype=np.float32),
)
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)
def smooth_depth_map(depth, smoothing):
import numpy as np
factor = min(max(float(smoothing) / 100.0, 0.0), 1.0)
if factor <= 0.0:
return np.asarray(depth, dtype=np.float32).copy()
import cv2
values = np.asarray(depth, dtype=np.float32)
valid = np.isfinite(values) & (values > 0.0)
if not valid.any():
raise ValueError("DA3 returned no valid depth values")
low, high = np.percentile(values[valid], (5.0, 95.0))
depth_range = max(float(high - low), 1e-6)
diameter = 3 + 2 * int(round(3.0 * factor))
working_values = values.copy()
working_values[~valid] = float(np.median(values[valid]))
smoothed = cv2.bilateralFilter(
working_values,
diameter,
depth_range * (0.01 + 0.09 * factor),
1.0 + 4.0 * factor,
)
smoothed[~valid] = values[~valid]
return smoothed
def compact_mesh(vertices, faces):
import numpy as np
vertices = np.asarray(vertices, dtype=np.float32)
faces = np.asarray(faces, dtype=np.int32)
used_vertices = np.unique(faces)
index_map = np.full(len(vertices), -1, dtype=np.int32)
index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32)
vertices = vertices[used_vertices]
faces = index_map[faces]
return vertices, faces
def project_camera_uv(vertices, intrinsics, image_size):
import numpy as np
width, height = image_size
z = vertices[:, 2]
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
raise ValueError("Simplified mesh contains invalid camera depth")
u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2]
v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2]
denominator_x = max(width - 1, 1)
denominator_y = max(height - 1, 1)
return np.column_stack(
(
np.clip(u / denominator_x, 0.0, 1.0),
np.clip(1.0 - v / denominator_y, 0.0, 1.0),
)
).astype(np.float32)
def project_camera_pixels(vertices, intrinsics):
import numpy as np
z = vertices[:, 2]
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
raise ValueError("Mesh contains invalid camera depth")
return np.column_stack(
(
intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2],
intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2],
)
)
def sample_image_nearest(values, pixels):
import numpy as np
height, width = values.shape
x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1)
y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1)
return np.asarray(values)[y, x]
def confidence_percentile_map(depth, confidence):
import numpy as np
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
result = np.zeros(depth.shape, dtype=np.float32)
if not valid.any():
return result
sorted_values = np.sort(np.asarray(confidence)[valid])
ranks = np.searchsorted(
sorted_values,
np.asarray(confidence)[valid],
side="right",
)
result[valid] = ranks.astype(np.float32) / len(sorted_values)
return result
def triangle_edge_stretch_values(points, triangles):
import numpy as np
flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3)
triangle_points = flat_points[triangles]
triangle_depth = triangle_points[..., 2]
safe_depth = np.where(
np.abs(triangle_depth) > 1e-8,
triangle_depth,
np.nan,
)
triangle_rays = triangle_points / safe_depth[..., None]
next_points = np.roll(triangle_points, -1, axis=1)
next_rays = np.roll(triangle_rays, -1, axis=1)
next_depth = np.roll(triangle_depth, -1, axis=1)
local_depth = (triangle_depth + next_depth) * 0.5
flat_edges = (
triangle_rays * local_depth[..., None]
- next_rays * local_depth[..., None]
)
flat_lengths = np.linalg.norm(flat_edges, axis=-1)
actual_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
stretch = actual_lengths / np.maximum(flat_lengths, 1e-8)
stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0
return stretch
def image_mesh_attributes(
vertices,
faces,
intrinsics,
depth,
confidence,
confidence_percentiles,
foreground_mask,
):
import numpy as np
pixels = project_camera_pixels(vertices, intrinsics)
point_confidence = sample_image_nearest(confidence, pixels).astype(
np.float32
)
point_confidence = np.nan_to_num(point_confidence, nan=0.0)
point_percentile = sample_image_nearest(
confidence_percentiles,
pixels,
).astype(np.float32)
point_background = ~sample_image_nearest(foreground_mask, pixels).astype(
bool
)
face_confidence = point_percentile[faces].min(axis=1)
face_edge_stretch = triangle_edge_stretch_values(
vertices,
faces,
).max(axis=1)
face_background = point_background[faces].any(axis=1)
return {
"point_depth": vertices[:, 2].astype(np.float32),
"point_confidence": point_confidence,
"point_confidence_percentile": point_percentile,
"point_background": point_background,
"face_confidence_percentile": face_confidence.astype(np.float32),
"face_edge_stretch": face_edge_stretch.astype(np.float32),
"face_background": face_background,
}
def export_prediction_attributes(prediction, output_directory):
import numpy as np
values = {}
for name in ("depth", "conf", "intrinsics", "extrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.asarray(value)
if "depth" in values and "conf" in values:
values["confidence_percentile"] = np.stack(
[
confidence_percentile_map(depth, confidence)
for depth, confidence in zip(
values["depth"],
values["conf"],
)
]
)
output_path = Path(output_directory) / "prediction.attributes.npz"
np.savez(output_path, **values)
return output_path
def export_image_mesh_data(
prediction,
image_path,
output_directory,
foreground_path=None,
progress_callback=None,
):
import numpy as np
from PIL import Image
total_steps = 6
def progress(step, label):
if progress_callback is not None:
progress_callback(step, total_steps, label)
progress(1, "Validating DA3 data")
required = ("depth", "conf", "intrinsics")
missing = [
name
for name in required
if getattr(prediction, name, None) is None
]
if missing:
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
depth_frames = np.asarray(prediction.depth, dtype=np.float32)
confidence_frames = np.asarray(prediction.conf, dtype=np.float32)
if depth_frames.shape[0] != 1:
raise ValueError("Convert to Mesh currently supports one image")
raw_depth = depth_frames[0]
confidence = confidence_frames[0]
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
progress(2, "Unprojecting camera geometry")
points = unproject_camera_depth(raw_depth, intrinsics)
progress(3, "Preparing DA3 and BEN2 attributes")
foreground_mask = np.ones(raw_depth.shape, dtype=bool)
if foreground_path is not None:
foreground_mask = foreground_mesh_mask(
foreground_path,
raw_depth.shape,
)
confidence_percentiles = confidence_percentile_map(
raw_depth,
confidence,
)
progress(4, "Building editable source topology")
faces = mesh_triangle_faces(
points,
raw_depth,
confidence,
None,
False,
0.0,
)
if not len(faces):
raise ValueError(
"DA3 output does not contain any valid mesh faces"
)
progress(5, "Projecting UV and structured attributes")
vertices, faces = compact_mesh(
points.reshape(-1, 3),
faces,
)
height, width = raw_depth.shape
uv = project_camera_uv(vertices, intrinsics, (width, height))
attributes = image_mesh_attributes(
vertices,
faces,
intrinsics,
raw_depth,
confidence,
confidence_percentiles,
foreground_mask,
)
valid_depth = raw_depth[np.isfinite(raw_depth) & (raw_depth > 0.0)]
reference_depth = float(np.median(valid_depth))
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
progress(6, "Writing mesh data")
result_path = output_directory / "mesh.npz"
np.savez_compressed(
result_path,
vertices=vertices,
faces=faces,
uv=uv,
intrinsics=intrinsics,
image_size=np.asarray((width, height), dtype=np.int32),
reference_depth=np.asarray(reference_depth, dtype=np.float32),
**attributes,
)
texture_path = output_directory / "texture.png"
with Image.open(image_path) as image:
image.convert("RGBA").save(texture_path, format="PNG", compress_level=1)
return result_path, texture_path
def depth_edge_mask(depth, threshold=12.0 / 255.0):
import numpy as np
padded = np.pad(depth.astype(np.float32), 1, mode="edge")
grad_x = (
-padded[:-2, :-2]
+ padded[:-2, 2:]
- 2.0 * padded[1:-1, :-2]
+ 2.0 * padded[1:-1, 2:]
- padded[2:, :-2]
+ padded[2:, 2:]
)
grad_y = (
-padded[:-2, :-2]
- 2.0 * padded[:-2, 1:-1]
- padded[:-2, 2:]
+ padded[2:, :-2]
+ 2.0 * padded[2:, 1:-1]
+ padded[2:, 2:]
)
magnitude = np.hypot(grad_x, grad_y)
finite = np.isfinite(magnitude)
if not finite.any():
return np.ones(depth.shape, dtype=bool)
minimum = float(magnitude[finite].min())
maximum = float(magnitude[finite].max())
if maximum <= minimum:
return np.zeros(depth.shape, dtype=bool)
normalized = (magnitude - minimum) / (maximum - minimum)
return normalized >= threshold
values = np.asarray(depth, dtype=np.float32)
edges = np.zeros(values.shape, dtype=bool)
for axis in (0, 1):
first = np.take(values, range(values.shape[axis] - 1), axis=axis)
second = np.take(values, range(1, values.shape[axis]), axis=axis)
valid = (
np.isfinite(first)
& np.isfinite(second)
& (first > 0.0)
& (second > 0.0)
)
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
discontinuity = valid & (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 edge_culling_stretch_limit(culling):
factor = min(max(float(culling) / 100.0, 0.0001), 1.0)
return 1.0 / factor
def foreground_mesh_mask(foreground_path, depth_shape):
import numpy as np
from PIL import Image, ImageFilter
height, width = depth_shape
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize(
(width, height),
resample=Image.Resampling.BILINEAR,
)
alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75))
return np.asarray(alpha) >= 128
def confidence_percentile_threshold(depth, confidence, percentile):
import numpy as np
if float(percentile) <= 0.0:
return None
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
if not valid.any():
raise ValueError("DA3 returned no valid depth confidence values")
@@ -416,6 +739,9 @@ def mesh_triangle_faces(
confidence_threshold,
filter_edges,
max_edge_length,
edge_threshold=12.0 / 255.0,
edge_stretch_limit=None,
vertex_mask=None,
):
import numpy as np
@@ -434,24 +760,47 @@ def mesh_triangle_faces(
np.isfinite(points).all(axis=-1)
& np.isfinite(depth)
& (depth > 0.0)
& (confidence >= confidence_threshold)
)
if confidence_threshold is not None:
valid &= np.isfinite(confidence) & (
confidence >= confidence_threshold
)
if filter_edges:
valid &= ~depth_edge_mask(depth)
face_valid = valid.reshape(-1)[quads].all(axis=1)
if max_edge_length > 0.0:
flat_points = points.reshape(-1, 3)
quad_points = flat_points[quads]
next_points = np.roll(quad_points, -1, axis=1)
edge_lengths = np.linalg.norm(quad_points - next_points, axis=-1)
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
quads = quads[face_valid]
if not len(quads):
return np.empty((0, 3), dtype=np.int32)
return np.concatenate(
valid &= ~depth_edge_mask(depth, edge_threshold)
if vertex_mask is not None:
if np.shape(vertex_mask) != depth.shape:
raise ValueError("Foreground mask and depth dimensions do not match")
valid &= np.asarray(vertex_mask, dtype=bool)
triangles = np.concatenate(
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
axis=0,
)
face_valid = valid.reshape(-1)[triangles].all(axis=1)
if edge_stretch_limit is not None:
face_valid &= triangle_edge_stretch_mask(
points,
triangles,
edge_stretch_limit,
)
if max_edge_length > 0.0:
flat_points = points.reshape(-1, 3)
triangle_points = flat_points[triangles]
next_points = np.roll(triangle_points, -1, axis=1)
edge_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
triangles = triangles[face_valid]
if not len(triangles):
return np.empty((0, 3), dtype=np.int32)
return triangles
def triangle_edge_stretch_mask(points, triangles, stretch_limit):
stretch = triangle_edge_stretch_values(points, triangles)
tolerance = float(stretch_limit) * (1.0 + 1e-5)
return (stretch <= tolerance).all(axis=1)
def export_textured_mesh_glb(
@@ -512,6 +861,16 @@ def export_textured_mesh_glb(
visual=visual,
process=False,
)
mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1)
mesh.vertex_attributes["o_confidence"] = confidence[
index
].reshape(-1)
mesh.vertex_attributes["o_confidence_percentile"] = (
confidence_percentile_map(
depth[index],
confidence[index],
).reshape(-1)
)
mesh.remove_unreferenced_vertices()
meshes.append(mesh)
referenced_points.append(mesh.vertices)
@@ -1160,6 +1519,9 @@ def exported_paths(output_directory, export_format):
if export_format == "mesh":
result_path = output_directory / "mesh.glb"
return [result_path] if result_path.is_file() else []
if export_format == "mesh_data":
result_path = output_directory / "mesh.npz"
return [result_path] if result_path.is_file() else []
if export_format == "mini_npz":
result_path = output_directory / "results.npz"
if result_path.is_file():
@@ -1401,6 +1763,23 @@ def reconstruct(args, model_cache=None, cancel_check=None):
bool(getattr(args, "mesh_filter_edges", True)),
float(getattr(args, "mesh_max_edge_length", 0.1)),
)
elif args.export_format == "mesh_data":
def mesh_progress(step, total, label):
write_status(
args.status,
0.88 + 0.09 * step / total,
f"Post-processing {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,
)
else:
export(
prediction,
@@ -1408,6 +1787,13 @@ def reconstruct(args, model_cache=None, cancel_check=None):
str(output_directory),
**export_kwargs,
)
if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
write_status(
args.status,
0.97,
"Writing structured debug attributes",
)
export_prediction_attributes(prediction, output_directory)
check_cancelled(cancel_check)
result_paths = wait_for_exported_paths(