nan
07d836d9ad
- 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.
565 lines
18 KiB
Python
565 lines
18 KiB
Python
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"}
|