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

Implement Balloon Depth Sampling and Validation Tools

Browse Source 
- Added depth.py for MoGe-2 depth sampling and balloon deformation functionality.
- Introduced validate_balloon_asset.py to validate balloon shape modifiers and material properties in Blender.
- Enhanced test_server.py with new unit tests for balloon depth analysis and job progress tracking.
- Updated existing tests to improve code clarity and maintainability.
This commit is contained in:
马南
2026-08-21 08:27:53 +08:00
parent 8afcf48ddf
commit b5a2f9690c
17 changed files with 1952 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/snapform/operators/__init__.py
View File

@@ -2,7 +2,7 @@ from .background import ProcessImageBackground, RemoveImageBackground
from .balloon import (
ActivateBalloonTool,
ExtractBalloon,
RefineBalloon,
AnalyzeBalloon,
)
from .debug import DebugBen2, DebugDa3, DebugMoge2
from .depth_plane import (
@@ -31,7 +31,7 @@ CLASSES = (
ExtractOriginalFrames,
ClearModels,
ExtractBalloon,
RefineBalloon,
AnalyzeBalloon,
ActivateBalloonTool,
RemoveImageBackground,
ConvertToDepthPlane,

2
src/snapform/operators/balloon/__init__.py
View File

@@ -2,5 +2,5 @@ from .operator import (
ActivateBalloonTool,
BalloonTool,
ExtractBalloon,
RefineBalloon,
AnalyzeBalloon,
)

622
src/snapform/operators/balloon/depth.py Normal file
View File

@@ -0,0 +1,622 @@
"""MoGe-2 depth sampling and adjustable Balloon deformation."""
from pathlib import Path
import bpy
BALLOON_DEPTH_ATTRIBUTE = "o_balloon_depth"
BALLOON_SHAPE_NODE_GROUP_NAME = "O Balloon Shape"
BALLOON_SHAPE_NODE_GROUP_VERSION = 16
BALLOON_SHAPE_VERSION_PROPERTY = "snapform_balloon_shape_version"
def load_balloon_depth_data(path):
import numpy as np
try:
with np.load(Path(path), allow_pickle=False) as archive:
required = {
"depth_signal",
"baseline_vector",
"valid",
"origin",
"reference_depth",
"crop_size",
"source_size",
}
missing = required.difference(archive.files)
if missing:
raise ValueError(
f"Balloon depth data is missing: {', '.join(sorted(missing))}"
)
data = {
"depth_signal": np.asarray(
archive["depth_signal"],
dtype=np.float32,
),
"baseline_vector": np.asarray(
archive["baseline_vector"],
dtype=np.float32,
),
"valid": np.asarray(archive["valid"], dtype=bool),
"origin": tuple(int(value) for value in archive["origin"]),
"reference_depth": float(archive["reference_depth"]),
"crop_size": tuple(int(value) for value in archive["crop_size"]),
"source_size": tuple(int(value) for value in archive["source_size"]),
}
except OSError as error:
raise ValueError(f"Unable to read Balloon depth data: {path}") from error
if data["depth_signal"].ndim != 2:
raise ValueError("Balloon depth signal must be two-dimensional")
if data["valid"].shape != data["depth_signal"].shape:
raise ValueError("Balloon depth mask dimensions do not match")
if data["baseline_vector"].shape != (3,):
raise ValueError("Balloon base plane vector must have three components")
if any(value <= 0 for value in (*data["crop_size"], *data["source_size"])):
raise ValueError("Balloon depth image dimensions must be positive")
return data
def _sample_valid_bilinear(values, valid, x, y):
import numpy as np
height, width = values.shape
x = np.clip(x, 0.0, max(width - 1, 0))
y = np.clip(y, 0.0, max(height - 1, 0))
x0 = np.floor(x).astype(np.int64)
y0 = np.floor(y).astype(np.int64)
x1 = np.minimum(x0 + 1, width - 1)
y1 = np.minimum(y0 + 1, height - 1)
fx = x - x0
fy = y - y0
samples = np.stack(
(
values[y0, x0],
values[y0, x1],
values[y1, x0],
values[y1, x1],
),
axis=1,
)
weights = np.stack(
(
(1.0 - fx) * (1.0 - fy),
fx * (1.0 - fy),
(1.0 - fx) * fy,
fx * fy,
),
axis=1,
)
sample_valid = np.stack(
(
valid[y0, x0],
valid[y0, x1],
valid[y1, x0],
valid[y1, x1],
),
axis=1,
)
weights *= sample_valid
totals = weights.sum(axis=1)
result = np.zeros(len(x), dtype=np.float32)
usable = totals > 1e-8
result[usable] = (samples[usable] * weights[usable]).sum(axis=1) / totals[usable]
return result, usable
def balloon_depth_fields(vertices, vertex_uv, depth_data):
"""Sample camera depth and convert the fitted base plane to local space."""
import numpy as np
vertices = np.asarray(vertices, dtype=np.float32)
uv = np.asarray(vertex_uv, dtype=np.float32)
if vertices.ndim != 2 or vertices.shape[1] != 3:
raise ValueError("Balloon vertices must have three components")
if uv.shape != (len(vertices), 2):
raise ValueError("Balloon UV coordinates do not match its vertices")
depth_signal = np.asarray(depth_data["depth_signal"], dtype=np.float32)
valid = np.asarray(depth_data["valid"], dtype=bool)
origin_x, origin_y = depth_data["origin"]
crop_width, crop_height = depth_data["crop_size"]
source_width, source_height = depth_data["source_size"]
depth_height, depth_width = depth_signal.shape
crop_x = uv[:, 0] * source_width - origin_x
crop_y = (1.0 - uv[:, 1]) * source_height - origin_y
sample_x = crop_x * max(depth_width - 1, 0) / max(crop_width - 1, 1)
sample_y = crop_y * max(depth_height - 1, 0) / max(crop_height - 1, 1)
sampled, usable = _sample_valid_bilinear(
depth_signal,
valid,
sample_x,
sample_y,
)
inside_crop = (
(crop_x >= 0.0)
& (crop_x <= crop_width)
& (crop_y >= 0.0)
& (crop_y <= crop_height)
)
usable &= inside_crop
sampled[~usable] = 0.0
coordinates = np.stack(
(
sample_x / max(depth_width - 1, 1) * 2.0 - 1.0,
sample_y / max(depth_height - 1, 1) * 2.0 - 1.0,
np.ones(len(vertices), dtype=np.float32),
),
axis=1,
).astype(np.float32)
base_height = float(np.max(np.abs(vertices[:, 2])))
if not base_height > 0.0:
return np.zeros(len(vertices), dtype=np.float32), np.zeros(3, dtype=np.float32)
depth_scale = base_height
normalized_plane = np.asarray(
depth_data["baseline_vector"],
dtype=np.float32,
)
baseline_values = coordinates @ normalized_plane * depth_scale
local_design = np.stack(
(
vertices[:, 0],
vertices[:, 1],
np.ones(len(vertices), dtype=np.float32),
),
axis=1,
)
if np.count_nonzero(usable) >= 3:
local_plane = np.linalg.lstsq(
local_design[usable],
baseline_values[usable],
rcond=None,
)[0].astype(np.float32)
else:
local_plane = np.zeros(3, dtype=np.float32)
depth_values = sampled * depth_scale
if usable.any():
plane_without_offset = local_design[:, :2] @ local_plane[:2]
residual = depth_values - plane_without_offset
depth_values[usable] -= float(np.median(residual[usable]))
local_plane[2] = 0.0
return depth_values, local_plane
def _assign_point_attribute(mesh, name, values):
import numpy as np
values = np.asarray(values, dtype=np.float32)
if len(values) != len(mesh.vertices):
raise ValueError(f"Balloon attribute {name} does not match its vertices")
attribute = mesh.attributes.new(name=name, type="FLOAT", domain="POINT")
attribute.data.foreach_set("value", values)
def assign_balloon_shape_attributes(mesh, _vertices, depth_values=None):
if depth_values is not None:
_assign_point_attribute(mesh, BALLOON_DEPTH_ATTRIBUTE, depth_values)
def create_balloon_base_plane(context, balloon, vertices, local_plane):
import numpy as np
from mathutils import Vector
vertices = np.asarray(vertices, dtype=np.float32)
slope_y, slope_z, offset = (float(value) for value in local_plane)
extent = vertices[:, 1:].max(axis=0) - vertices[:, 1:].min(axis=0)
normal = Vector((1.0, -slope_y, -slope_z)).normalized()
controller = bpy.data.objects.new(f"{balloon.name} Base Plane", None)
context.collection.objects.link(controller)
controller.parent = balloon
controller.location = (0.0, 0.0, 0.0)
controller.rotation_euler = normal.to_track_quat("Z", "Y").to_euler()
controller.empty_display_type = "SINGLE_ARROW"
controller.empty_display_size = max(float(extent.max()) * 0.15, 0.01)
controller.show_in_front = True
controller.hide_render = True
controller["snapform_source_type"] = "BALLOON_BASE_PLANE"
controller["snapform_base_plane"] = (slope_y, slope_z, offset)
return controller
def _interface_socket(group, name, in_out, socket_type):
for item in group.interface.items_tree:
if item.item_type == "SOCKET" and item.name == name and item.in_out == in_out:
return item
return group.interface.new_socket(
name=name,
in_out=in_out,
socket_type=socket_type,
)
def _remove_obsolete_interface_sockets(group):
obsolete = {"Base X", "Base Y", "Base Offset", "Far Clamp"}
for item in list(group.interface.items_tree):
if (
item.item_type == "SOCKET"
and item.in_out == "INPUT"
and item.name in obsolete
):
group.interface.remove(item)
def _float_input(group, name, default, minimum, maximum, description):
socket = _interface_socket(group, name, "INPUT", "NodeSocketFloat")
socket.default_value = default
socket.min_value = minimum
socket.max_value = maximum
socket.description = description
if hasattr(socket, "structure_type"):
socket.structure_type = "SINGLE"
return socket
def _bool_input(group, name, default, description):
socket = _interface_socket(group, name, "INPUT", "NodeSocketBool")
socket.default_value = default
socket.description = description
if hasattr(socket, "structure_type"):
socket.structure_type = "SINGLE"
return socket
def balloon_shape_node_group():
group = bpy.data.node_groups.get(BALLOON_SHAPE_NODE_GROUP_NAME)
if (
group is not None
and group.get(BALLOON_SHAPE_VERSION_PROPERTY)
== BALLOON_SHAPE_NODE_GROUP_VERSION
):
return group
if group is None:
group = bpy.data.node_groups.new(
BALLOON_SHAPE_NODE_GROUP_NAME,
"GeometryNodeTree",
)
else:
group.nodes.clear()
_remove_obsolete_interface_sockets(group)
group.is_modifier = True
geometry_input = _interface_socket(
group,
"Geometry",
"INPUT",
"NodeSocketGeometry",
)
if hasattr(geometry_input, "structure_type"):
geometry_input.structure_type = "SINGLE"
_float_input(
group,
"Thickness",
1.0,
0.0,
10.0,
"Scale the symmetric base Balloon thickness along local X.",
)
_float_input(
group,
"Depth Strength",
1.0,
0.0,
10.0,
"Scale the signed MoGe-2 depth variation.",
)
_bool_input(
group,
"Double Sided Depth",
False,
"Apply AI depth to both sides; keep disabled for a stable Balloon back.",
)
base_plane_input = _interface_socket(
group,
"Base Plane",
"INPUT",
"NodeSocketObject",
)
base_plane_input.description = (
"Rotate the child Empty for middle-plane orientation; move it along local Z "
"to set the neutral depth level."
)
_interface_socket(group, "Geometry", "OUTPUT", "NodeSocketGeometry")
nodes = group.nodes
links = group.links
group_input = nodes.new("NodeGroupInput")
group_input.location = (-1500.0, 500.0)
position = nodes.new("GeometryNodeInputPosition")
position.location = (-1500.0, 200.0)
separate = nodes.new("ShaderNodeSeparateXYZ")
separate.location = (-1300.0, 200.0)
base_scale = nodes.new("ShaderNodeMath")
base_scale.operation = "MULTIPLY"
base_scale.location = (-1050.0, 350.0)
base_magnitude = nodes.new("ShaderNodeMath")
base_magnitude.operation = "ABSOLUTE"
base_magnitude.location = (-800.0, 350.0)
side = nodes.new("ShaderNodeMath")
side.operation = "SIGN"
side.location = (-1050.0, 100.0)
front = nodes.new("ShaderNodeMath")
front.operation = "LESS_THAN"
front.inputs[1].default_value = 0.0
front.location = (750.0, -350.0)
depth_attribute = nodes.new("GeometryNodeInputNamedAttribute")
depth_attribute.data_type = "FLOAT"
depth_attribute.inputs["Name"].default_value = BALLOON_DEPTH_ATTRIBUTE
depth_attribute.location = (-1500.0, -150.0)
profile_factor = nodes.new("ShaderNodeMath")
profile_factor.operation = "ABSOLUTE"
profile_factor.location = (-800.0, 100.0)
profile_statistics = nodes.new("GeometryNodeAttributeStatistic")
profile_statistics.data_type = "FLOAT"
profile_statistics.domain = "POINT"
profile_statistics.location = (-600.0, 50.0)
profile_max_safe = nodes.new("ShaderNodeMath")
profile_max_safe.operation = "MAXIMUM"
profile_max_safe.inputs[1].default_value = 1e-6
profile_max_safe.location = (-350.0, 100.0)
profile_normalized = nodes.new("ShaderNodeMath")
profile_normalized.operation = "DIVIDE"
profile_normalized.location = (-100.0, 100.0)
profile_smooth = nodes.new("ShaderNodeMapRange")
profile_smooth.data_type = "FLOAT"
profile_smooth.interpolation_type = "SMOOTHERSTEP"
profile_smooth.clamp = True
profile_smooth.inputs["From Min"].default_value = 0.0
profile_smooth.inputs["From Max"].default_value = 1.0
profile_smooth.inputs["To Min"].default_value = 0.0
profile_smooth.inputs["To Max"].default_value = 1.0
profile_smooth.location = (100.0, 50.0)
profile_active = nodes.new("ShaderNodeMath")
profile_active.operation = "MULTIPLY"
profile_active.location = (350.0, 50.0)
object_info = nodes.new("GeometryNodeObjectInfo")
object_info.transform_space = "RELATIVE"
object_info.location = (-1500.0, -450.0)
plane_normal = nodes.new("FunctionNodeRotateVector")
plane_normal.inputs["Vector"].default_value = (0.0, 0.0, 1.0)
plane_normal.location = (-1250.0, -500.0)
effective_normal = nodes.new("GeometryNodeSwitch")
effective_normal.input_type = "VECTOR"
effective_normal.inputs["False"].default_value = (1.0, 0.0, 0.0)
effective_normal.location = (-1050.0, -500.0)
normal_separate = nodes.new("ShaderNodeSeparateXYZ")
normal_separate.location = (-1000.0, -550.0)
zero_position = nodes.new("ShaderNodeCombineXYZ")
zero_position.location = (-1050.0, -250.0)
plane_dot = nodes.new("ShaderNodeVectorMath")
plane_dot.operation = "DOT_PRODUCT"
plane_dot.location = (-800.0, -300.0)
plane_offset = nodes.new("ShaderNodeVectorMath")
plane_offset.operation = "DOT_PRODUCT"
plane_offset.location = (-800.0, -500.0)
plane_numerator = nodes.new("ShaderNodeMath")
plane_numerator.operation = "SUBTRACT"
plane_numerator.location = (-550.0, -300.0)
plane_divide = nodes.new("ShaderNodeMath")
plane_divide.operation = "DIVIDE"
plane_divide.location = (-300.0, -300.0)
depth_residual = nodes.new("ShaderNodeMath")
depth_residual.operation = "ADD"
depth_residual.location = (-50.0, -200.0)
depth_calibration = nodes.new("ShaderNodeMath")
depth_calibration.operation = "MULTIPLY"
depth_calibration.inputs[1].default_value = 0.5
depth_calibration.label = "Calibrate Relative Depth"
depth_calibration.location = (200.0, -200.0)
corrected_depth = nodes.new("ShaderNodeMath")
corrected_depth.operation = "ADD"
corrected_depth.location = (450.0, -200.0)
depth_scale = nodes.new("ShaderNodeMath")
depth_scale.operation = "MULTIPLY"
depth_scale.location = (700.0, -150.0)
depth_mix = nodes.new("ShaderNodeMix")
depth_mix.data_type = "FLOAT"
depth_mix.clamp_factor = True
depth_mix.location = (450.0, -50.0)
signed_depth = nodes.new("ShaderNodeMath")
signed_depth.operation = "MULTIPLY"
signed_depth.location = (850.0, 50.0)
front_choice = nodes.new("GeometryNodeSwitch")
front_choice.input_type = "FLOAT"
front_choice.location = (1100.0, 0.0)
double_choice = nodes.new("GeometryNodeSwitch")
double_choice.input_type = "FLOAT"
double_choice.location = (1350.0, 100.0)
combine = nodes.new("ShaderNodeCombineXYZ")
combine.location = (1600.0, 250.0)
set_position = nodes.new("GeometryNodeSetPosition")
set_position.location = (1850.0, 450.0)
group_output = nodes.new("NodeGroupOutput")
group_output.location = (2450.0, 450.0)
base_frame = nodes.new("NodeFrame")
base_frame.label = "Base Thickness"
plane_frame = nodes.new("NodeFrame")
plane_frame.label = "Middle Plane"
mix_frame = nodes.new("NodeFrame")
mix_frame.label = "Depth Mix"
output_frame = nodes.new("NodeFrame")
output_frame.label = "Side Output"
for node in (
separate,
base_scale,
base_magnitude,
profile_factor,
profile_statistics,
profile_max_safe,
profile_normalized,
profile_smooth,
):
node.parent = base_frame
for node in (
object_info,
plane_normal,
effective_normal,
normal_separate,
zero_position,
plane_dot,
plane_offset,
plane_numerator,
plane_divide,
):
node.parent = plane_frame
for node in (
depth_residual,
depth_calibration,
corrected_depth,
depth_scale,
profile_active,
depth_mix,
):
node.parent = mix_frame
for node in (
side,
front,
signed_depth,
front_choice,
double_choice,
combine,
set_position,
):
node.parent = output_frame
base_frame.location = (-1800.0, 500.0)
separate.location = (0.0, 200.0)
base_scale.location = (220.0, 350.0)
base_magnitude.location = (440.0, 350.0)
profile_factor.location = (220.0, 50.0)
profile_statistics.location = (440.0, 0.0)
profile_max_safe.location = (660.0, 0.0)
profile_normalized.location = (880.0, 0.0)
profile_smooth.location = (1100.0, 0.0)
plane_frame.location = (-1800.0, -600.0)
object_info.location = (0.0, -300.0)
plane_normal.location = (220.0, -300.0)
effective_normal.location = (440.0, -300.0)
normal_separate.location = (660.0, -300.0)
zero_position.location = (220.0, 0.0)
plane_dot.location = (440.0, 0.0)
plane_offset.location = (660.0, -100.0)
plane_numerator.location = (880.0, 0.0)
plane_divide.location = (1100.0, 0.0)
mix_frame.location = (-200.0, 200.0)
profile_active.location = (0.0, 300.0)
depth_residual.location = (0.0, 0.0)
depth_calibration.location = (220.0, 0.0)
corrected_depth.location = (440.0, 0.0)
depth_scale.location = (660.0, 0.0)
depth_mix.location = (880.0, 150.0)
output_frame.location = (900.0, 200.0)
side.location = (0.0, 0.0)
front.location = (0.0, -250.0)
signed_depth.location = (220.0, 200.0)
front_choice.location = (440.0, 100.0)
double_choice.location = (660.0, 150.0)
combine.location = (880.0, 150.0)
set_position.location = (1100.0, 150.0)
group_input.location = (-2200.0, 500.0)
position.location = (-2200.0, 200.0)
depth_attribute.location = (-2200.0, -200.0)
group_output.location = (2500.0, 350.0)
links.new(position.outputs["Position"], separate.inputs["Vector"])
links.new(separate.outputs["X"], base_scale.inputs[0])
links.new(group_input.outputs["Thickness"], base_scale.inputs[1])
links.new(base_scale.outputs[0], base_magnitude.inputs[0])
links.new(separate.outputs["X"], side.inputs[0])
links.new(separate.outputs["X"], front.inputs[0])
links.new(separate.outputs["X"], profile_factor.inputs[0])
links.new(group_input.outputs["Geometry"], profile_statistics.inputs["Geometry"])
links.new(profile_factor.outputs[0], profile_statistics.inputs["Attribute"])
links.new(profile_statistics.outputs["Max"], profile_max_safe.inputs[0])
links.new(profile_factor.outputs[0], profile_normalized.inputs[0])
links.new(profile_max_safe.outputs[0], profile_normalized.inputs[1])
links.new(profile_normalized.outputs[0], profile_smooth.inputs["Value"])
links.new(profile_smooth.outputs["Result"], profile_active.inputs[0])
links.new(depth_attribute.outputs["Exists"], profile_active.inputs[1])
links.new(group_input.outputs["Base Plane"], object_info.inputs["Object"])
links.new(object_info.outputs["Rotation"], plane_normal.inputs["Rotation"])
links.new(depth_attribute.outputs["Exists"], effective_normal.inputs["Switch"])
links.new(plane_normal.outputs["Vector"], effective_normal.inputs["True"])
links.new(effective_normal.outputs["Output"], normal_separate.inputs["Vector"])
links.new(separate.outputs["Y"], zero_position.inputs["Y"])
links.new(separate.outputs["Z"], zero_position.inputs["Z"])
links.new(zero_position.outputs["Vector"], plane_dot.inputs[0])
links.new(effective_normal.outputs["Output"], plane_dot.inputs[1])
links.new(object_info.outputs["Location"], plane_offset.inputs[0])
links.new(effective_normal.outputs["Output"], plane_offset.inputs[1])
links.new(plane_offset.outputs["Value"], plane_numerator.inputs[0])
links.new(plane_dot.outputs["Value"], plane_numerator.inputs[1])
links.new(plane_numerator.outputs[0], plane_divide.inputs[0])
links.new(normal_separate.outputs["X"], plane_divide.inputs[1])
links.new(depth_attribute.outputs["Attribute"], depth_residual.inputs[0])
links.new(plane_divide.outputs[0], depth_residual.inputs[1])
links.new(depth_residual.outputs[0], depth_calibration.inputs[0])
links.new(depth_calibration.outputs[0], corrected_depth.inputs[0])
links.new(profile_statistics.outputs["Max"], corrected_depth.inputs[1])
links.new(corrected_depth.outputs[0], depth_scale.inputs[0])
links.new(group_input.outputs["Depth Strength"], depth_scale.inputs[1])
links.new(profile_active.outputs[0], depth_mix.inputs["Factor"])
links.new(base_magnitude.outputs[0], depth_mix.inputs["A"])
links.new(depth_scale.outputs[0], depth_mix.inputs["B"])
links.new(side.outputs[0], signed_depth.inputs[0])
links.new(depth_mix.outputs["Result"], signed_depth.inputs[1])
links.new(front.outputs[0], front_choice.inputs["Switch"])
links.new(base_scale.outputs[0], front_choice.inputs["False"])
links.new(signed_depth.outputs[0], front_choice.inputs["True"])
links.new(
group_input.outputs["Double Sided Depth"],
double_choice.inputs["Switch"],
)
links.new(front_choice.outputs["Output"], double_choice.inputs["False"])
links.new(signed_depth.outputs[0], double_choice.inputs["True"])
links.new(double_choice.outputs["Output"], combine.inputs["X"])
links.new(separate.outputs["Y"], combine.inputs["Y"])
links.new(separate.outputs["Z"], combine.inputs["Z"])
links.new(group_input.outputs["Geometry"], set_position.inputs["Geometry"])
links.new(combine.outputs["Vector"], set_position.inputs["Position"])
links.new(set_position.outputs["Geometry"], group_output.inputs["Geometry"])
group.color_tag = "GEOMETRY"
group.description = (
"Mix Balloon thickness with depth around an adjustable child Empty middle plane."
)
group[BALLOON_SHAPE_VERSION_PROPERTY] = BALLOON_SHAPE_NODE_GROUP_VERSION
group.use_fake_user = True
return group
def add_balloon_shape_modifier(
balloon,
thickness,
base_plane=None,
):
from ..material import modifier_input_identifier, set_modifier_input
node_group = balloon_shape_node_group()
modifier = balloon.modifiers.new(name=node_group.name, type="NODES")
modifier.node_group = node_group
modifier.show_group_selector = False
values = {
"Thickness": thickness,
"Depth Strength": 1.0,
"Double Sided Depth": False,
"Base Plane": base_plane,
}
for name, value in values.items():
identifier = modifier_input_identifier(node_group, name)
set_modifier_input(modifier, identifier, value)
return modifier

293
src/snapform/operators/balloon/operator.py
View File

@@ -3,10 +3,19 @@ import math
import bpy
from .depth import (
add_balloon_shape_modifier,
assign_balloon_shape_attributes,
balloon_depth_fields,
create_balloon_base_plane,
load_balloon_depth_data,
)
from .geometry import INFLATION, POINT_SPACING, triangulate_balloon
from ...properties import (
BEN2_MODEL_KEY,
balloon_ai_refine_property as ai_refine_property,
balloon_contour_analysis_property as contour_analysis_property,
balloon_depth_analysis_property as depth_analysis_property,
balloon_depth_model_property as depth_model_property,
balloon_mode_property as mode_property,
balloon_resolution_property as resolution_property,
balloon_reuse_material_property as reuse_material_property,
@@ -25,16 +34,24 @@ from ..image import (
color_data_name,
is_image_empty,
material_data_name,
normal_data_name,
plane_bounds,
prepare_image_input,
remove_temporary_input,
video_total_frames,
)
from ..material import TEXTURE_MATERIAL_NODE_GROUP_NAME, create_material
from ..material import (
BALLOON_MATERIAL_NODE_GROUP_NAME,
TEXTURE_MATERIAL_NODE_GROUP_NAME,
create_material,
)
from ..surface import assign_vertex_uv
MIN_PATH_DISTANCE = 2.0
BALLOON_MATERIAL_PROPERTY = "snapform_balloon_material"
BALLOON_SOURCE_IMAGE_PROPERTY = "snapform_source_image"
BALLOON_NORMAL_IMAGE_PROPERTY = "snapform_balloon_normal_image"
def simplify_screen_path(path):
@@ -150,7 +167,45 @@ def material_uses_balloon_texture(material, color_name):
return has_texture_layer and has_color_image
def reusable_balloon_material(material_name, color_name):
def marked_balloon_material(source_image, normal_required=False):
return next(
(
material
for material in bpy.data.materials
if callable(getattr(material, "get", None))
and bool(material.get(BALLOON_MATERIAL_PROPERTY))
and material.get(BALLOON_SOURCE_IMAGE_PROPERTY) is source_image
and (
material.get(BALLOON_NORMAL_IMAGE_PROPERTY) is not None
)
== normal_required
),
None,
)
def mark_balloon_material(material, source_image, normal_image=None):
material[BALLOON_MATERIAL_PROPERTY] = True
material[BALLOON_SOURCE_IMAGE_PROPERTY] = source_image
if normal_image is not None:
material[BALLOON_NORMAL_IMAGE_PROPERTY] = normal_image
def reusable_balloon_material(
material_name,
color_name,
source_image=None,
normal_required=False,
):
marked = (
marked_balloon_material(source_image, normal_required)
if source_image is not None
else None
)
if marked is not None:
return marked
if normal_required:
return None
return next(
(
material
@@ -162,6 +217,35 @@ def reusable_balloon_material(material_name, color_name):
)
def load_balloon_normal_image(path, source_object):
image = bpy.data.images.load(
str(path),
check_existing=False,
)
image.name = normal_data_name(source_object)
image.colorspace_settings.name = "Non-Color"
return image
def merge_balloon_normal_image(target, update):
import numpy as np
if tuple(target.size) != tuple(update.size):
raise ValueError("Reusable Balloon normal image dimensions do not match")
target_pixels = np.empty(len(target.pixels), dtype=np.float32)
update_pixels = np.empty(len(update.pixels), dtype=np.float32)
target.pixels.foreach_get(target_pixels)
update.pixels.foreach_get(update_pixels)
target_rgba = target_pixels.reshape((-1, 4))
update_rgba = update_pixels.reshape((-1, 4))
neutral = np.asarray((128.0 / 255.0, 128.0 / 255.0, 1.0))
changed = np.max(np.abs(update_rgba[:, :3] - neutral), axis=1) > (1.5 / 255.0)
target_rgba[changed] = update_rgba[changed]
target.pixels.foreach_set(target_pixels)
target.pack()
target.update()
def reusable_color_image(color_name):
return next(
(
@@ -173,6 +257,27 @@ def reusable_color_image(color_name):
)
BALLOON_AXIS_TO_SOURCE = (
(0.0, -1.0, 0.0, 0.0),
(0.0, 0.0, 1.0, 0.0),
(-1.0, 0.0, 0.0, 0.0),
(0.0, 0.0, 0.0, 1.0),
)
def orient_balloon_vertices(vertices):
"""Rotate legacy XY/Z Balloon coordinates into a YZ/-X frame."""
return [(-z, -x, y) for x, y, z in vertices]
def orient_balloon_world_matrix(source_matrix):
from mathutils import Matrix
return source_matrix @ Matrix(BALLOON_AXIS_TO_SOURCE)
def orient_balloon_base_plane(local_plane):
slope_x, slope_y, offset = local_plane
return (slope_x, -slope_y, -offset)
def extract_balloon_object(
context,
source_object,
@@ -183,6 +288,8 @@ def extract_balloon_object(
resolution=POINT_SPACING,
select_result=True,
reuse_material=True,
depth_data=None,
normal_path=None,
):
image = source_object.data
image_size = tuple(image.size)
@@ -199,45 +306,87 @@ def extract_balloon_object(
spacing=point_spacing,
world_metric=world_metric,
world_z_scale=world_z_scale,
inflation=inflation,
inflation=INFLATION,
algorithm=algorithm,
)
depth_values = None
local_base_plane = (0.0, 0.0, 0.0)
if depth_data is not None:
depth_values, local_base_plane = balloon_depth_fields(
vertices,
vertex_uv,
depth_data,
)
local_base_plane = orient_balloon_base_plane(local_base_plane)
vertices = orient_balloon_vertices(vertices)
source_name = source_object.name
balloon_name = source_name
material_name = material_data_name(source_object)
color_name = color_data_name(source_object)
color_image = None
uses_normal = normal_path is not None
material = (
reusable_balloon_material(material_name, color_name)
reusable_balloon_material(
material_name,
color_name,
image,
normal_required=uses_normal,
)
if reuse_material
else None
)
normal_image = None
owns_color_image = False
owns_normal_image = False
owns_material = False
mesh = None
balloon = None
base_plane = None
try:
if material is not None and not uses_normal:
mark_balloon_material(material, image)
if material is not None and uses_normal:
normal_image = material.get(BALLOON_NORMAL_IMAGE_PROPERTY)
update_normal = load_balloon_normal_image(normal_path, source_object)
try:
merge_balloon_normal_image(normal_image, update_normal)
finally:
bpy.data.images.remove(update_normal)
if material is None:
color_image = (
reusable_color_image(color_name)
if reuse_material
else None
)
color_image = reusable_color_image(color_name) if reuse_material else None
if color_image is None:
color_image = image.copy()
color_image.name = color_name
owns_color_image = True
if normal_path is not None:
normal_image = load_balloon_normal_image(
normal_path,
source_object,
)
normal_image.pack()
owns_normal_image = True
material = create_material(
material_name,
color_image,
0.0,
surface_group_name=TEXTURE_MATERIAL_NODE_GROUP_NAME,
1.0,
surface_group_name=(
BALLOON_MATERIAL_NODE_GROUP_NAME
if normal_image is not None
else TEXTURE_MATERIAL_NODE_GROUP_NAME
),
normal_image=normal_image,
scene=context.scene,
)
mark_balloon_material(material, image, normal_image)
owns_material = True
mesh = bpy.data.meshes.new(balloon_name)
mesh.from_pydata(vertices, [], faces)
assign_balloon_shape_attributes(
mesh,
vertices,
depth_values,
)
mesh.update()
assign_vertex_uv(mesh, vertex_uv)
for face in mesh.polygons:
@@ -245,26 +394,37 @@ def extract_balloon_object(
balloon = bpy.data.objects.new(balloon_name, mesh)
context.collection.objects.link(balloon)
balloon.matrix_world = matrix_world.copy()
balloon.matrix_world = orient_balloon_world_matrix(matrix_world)
balloon["snapform_source_type"] = "BALLOON"
balloon["snapform_balloon_algorithm"] = algorithm
balloon["snapform_source_object"] = source_name
balloon["snapform_inflation"] = inflation
balloon["snapform_resolution"] = point_spacing
balloon.data.materials.append(material)
if depth_data is not None:
base_plane = create_balloon_base_plane(
context,
balloon,
vertices,
local_base_plane,
)
add_balloon_shape_modifier(balloon, thickness, base_plane)
if depth_data is not None:
balloon["snapform_balloon_depth_analysis"] = True
balloon["snapform_reference_depth"] = depth_data["reference_depth"]
except Exception:
if base_plane is not None:
bpy.data.objects.remove(base_plane, do_unlink=True)
if balloon is not None:
bpy.data.objects.remove(balloon, do_unlink=True)
if mesh is not None and mesh.users == 0:
bpy.data.meshes.remove(mesh)
if owns_material and material is not None and material.users == 0:
bpy.data.materials.remove(material)
if (
owns_color_image
and color_image is not None
and color_image.users == 0
):
if owns_color_image and color_image is not None and color_image.users == 0:
bpy.data.images.remove(color_image)
if owns_normal_image and normal_image is not None and normal_image.users == 0:
bpy.data.images.remove(normal_image)
raise
bpy.ops.object.select_all(action="DESELECT")
@@ -323,7 +483,9 @@ class ExtractBalloon(bpy.types.Operator):
resolution: resolution_property()
source_object_name: hidden_string_property()
reuse_material: reuse_material_property()
ai_refine: ai_refine_property()
contour_analysis: contour_analysis_property()
depth_analysis: depth_analysis_property()
depth_model: depth_model_property()
source_matrix_data: hidden_string_property()
selection_data: hidden_string_property()
@@ -338,7 +500,10 @@ class ExtractBalloon(bpy.types.Operator):
layout.prop(self, "thickness")
layout.prop(self, "resolution", text="Resolution (m)")
layout.prop(self, "reuse_material")
layout.prop(self, "ai_refine")
layout.prop(self, "contour_analysis")
layout.prop(self, "depth_analysis")
if self.depth_analysis:
layout.prop(self, "depth_model")
@classmethod
def poll(cls, context):
@@ -349,9 +514,8 @@ class ExtractBalloon(bpy.types.Operator):
)
def _draw_lasso(self):
if (
not self._path
or not drawing_in_region(self._area_pointer, self._region_pointer)
if not self._path or not drawing_in_region(
self._area_pointer, self._region_pointer
):
return
vertices = list(self._path)
@@ -377,7 +541,9 @@ class ExtractBalloon(bpy.types.Operator):
self.thickness = settings.balloon_thickness
self.resolution = settings.balloon_resolution
self.reuse_material = settings.balloon_reuse_material
self.ai_refine = settings.balloon_ai_refine
self.contour_analysis = settings.balloon_contour_analysis
self.depth_analysis = settings.balloon_depth_analysis
self.depth_model = settings.balloon_depth_model
self.source_object_name = source_object.name
self.source_matrix_data = serialize_matrix(source_object.matrix_world)
self.selection_data = ""
@@ -449,14 +615,14 @@ class ExtractBalloon(bpy.types.Operator):
polygon = json.loads(self.selection_data)
if len(polygon) < 3:
raise ValueError("The stored lasso region is invalid")
if self.ai_refine:
if self.contour_analysis or self.depth_analysis:
if video_total_frames(source_object.data) > 1:
raise RuntimeError("AI Refine only supports a still image")
raise RuntimeError("AI Analysis only supports a still image")
source_path, delete_input = prepare_image_input(
source_object.data,
context.scene,
)
result = bpy.ops.snapform.refine_balloon(
result = bpy.ops.snapform.analyze_balloon(
input_path=str(source_path),
delete_input=delete_input,
source_object_name=self.source_object_name,
@@ -466,6 +632,9 @@ class ExtractBalloon(bpy.types.Operator):
thickness=self.thickness,
resolution=self.resolution,
reuse_material=self.reuse_material,
contour_analysis=self.contour_analysis,
depth_analysis=self.depth_analysis,
depth_model=self.depth_model,
)
if "CANCELLED" in result:
if delete_input:
@@ -495,12 +664,11 @@ class ExtractBalloon(bpy.types.Operator):
return {"FINISHED"}
class RefineBalloon(JobOperatorBase, bpy.types.Operator):
bl_idname = "snapform.refine_balloon"
bl_label = "Refine Balloon"
job_type = "refine-balloon"
starting_message = "Starting Balloon AI refinement"
class AnalyzeBalloon(JobOperatorBase, bpy.types.Operator):
bl_idname = "snapform.analyze_balloon"
bl_label = "Analyze Balloon"
job_type = "analyze-balloon"
starting_message = "Starting Balloon AI analysis"
input_path: bpy.props.StringProperty(options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
@@ -511,17 +679,35 @@ class RefineBalloon(JobOperatorBase, bpy.types.Operator):
thickness: thickness_property()
resolution: resolution_property()
reuse_material: reuse_material_property()
contour_analysis: bpy.props.BoolProperty(default=False, options=HIDDEN)
depth_analysis: bpy.props.BoolProperty(default=False, options=HIDDEN)
depth_model: depth_model_property()
resolution_level: bpy.props.IntProperty(
default=5,
min=0,
max=9,
options=HIDDEN,
)
def request(self, _context):
require_environment()
require_model(BEN2_MODEL_KEY)
if self.contour_analysis:
require_model(BEN2_MODEL_KEY)
if self.depth_analysis:
require_model(self.depth_model)
if not self.contour_analysis and not self.depth_analysis:
raise ValueError("Select Contour or Depth")
polygon = json.loads(self.selection_data)
if len(polygon) < 3:
raise ValueError("The stored lasso region is invalid")
return {
"input": str(self.input_path),
"model": BEN2_MODEL_KEY,
"contour_analysis": self.contour_analysis,
"depth_analysis": self.depth_analysis,
"contour_model": BEN2_MODEL_KEY,
"depth_model": self.depth_model,
"device": production_device().lower(),
"resolution_level": self.resolution_level,
"polygon": polygon,
}
@@ -529,12 +715,17 @@ class RefineBalloon(JobOperatorBase, bpy.types.Operator):
source_object = bpy.data.objects.get(self.source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source Image Empty is no longer available")
payload = json.loads(
result.file("contour").read_text(encoding="utf-8")
polygon = json.loads(self.selection_data)
if self.contour_analysis:
payload = json.loads(result.file("contour").read_text(encoding="utf-8"))
polygon = payload.get("points", ())
if len(polygon) < 3:
raise RuntimeError("BEN2 did not return a valid Balloon contour")
depth_data = (
load_balloon_depth_data(result.file("depth"))
if self.depth_analysis
else None
)
polygon = payload.get("points", ())
if len(polygon) < 3:
raise RuntimeError("BEN2 did not return a valid Balloon contour")
matrix_world = deserialize_matrix(self.source_matrix_data)
resolution_world = meters_to_world_distance(
context.scene,
@@ -550,15 +741,21 @@ class RefineBalloon(JobOperatorBase, bpy.types.Operator):
resolution=resolution_world,
select_result=False,
reuse_material=self.reuse_material,
depth_data=depth_data,
normal_path=(result.file("normal") if self.depth_analysis else None),
)
balloon["snapform_resolution_m"] = self.resolution
balloon["snapform_balloon_ai_refined"] = True
if self.contour_analysis:
balloon["snapform_balloon_contour_analysis"] = True
if self.depth_analysis:
balloon["snapform_balloon_depth_model"] = self.depth_model
push_undo("Extract Balloon")
return "AI-refined Balloon is ready"
return "AI-analyzed Balloon is ready"
def cleanup(self):
cleanup_temporary_input(self.input_path, self.delete_input)
class ActivateBalloonTool(bpy.types.Operator):
bl_idname = "snapform.activate_balloon_tool"
bl_label = "Balloon Lasso Tool"
@@ -610,7 +807,7 @@ class BalloonTool(bpy.types.WorkSpaceTool):
(
"zoom_path",
"scene.snapform_settings.balloon_thickness_radial_zoom",
)
),
]
},
),
@@ -626,10 +823,11 @@ class BalloonTool(bpy.types.WorkSpaceTool):
(
"zoom_path",
"scene.snapform_settings.balloon_resolution_radial_zoom",
)
),
]
},
), )
),
)
def draw_settings(context, layout, tool):
settings = context.scene.snapform_settings
@@ -637,4 +835,7 @@ class BalloonTool(bpy.types.WorkSpaceTool):
layout.prop(settings, "balloon_thickness")
layout.prop(settings, "balloon_resolution", text="Resolution (m)")
layout.prop(settings, "balloon_reuse_material")
layout.prop(settings, "balloon_ai_refine")
layout.prop(settings, "balloon_contour_analysis")
layout.prop(settings, "balloon_depth_analysis")
if settings.balloon_depth_analysis:
layout.prop(settings, "balloon_depth_model", text="Depth Model")

6
src/snapform/operators/balloon/poisson.py
View File

@@ -46,7 +46,11 @@ def cotangent_poisson_field(
edge_starts = []
edge_ends = []
weights = []
# 约束剖分的瘦三角形可能产生负 cotangent 权重；钳制后保持 Dirichlet 系统正定。
for (start, end), weight in edge_weights.items():
weight = max(float(weight), 0.0)
if weight <= 1e-14:
continue
start_compact = compact[start]
end_compact = compact[end]
if start_compact >= 0:
@@ -87,7 +91,7 @@ def cotangent_poisson_field(
break
multiplied = multiply(direction)
denominator = float(direction.dot(multiplied))
if denominator <= 1e-20 or not math.isfinite(denominator):
if denominator <= 0.0 or not math.isfinite(denominator):
raise ValueError("Unable to solve the Poisson Balloon field")
step = residual_product / denominator
solution += step * direction

19
src/snapform/operators/material.py
View File

@@ -17,6 +17,7 @@ DEPTH_SCALE_ATTRIBUTE = "o_depth_scale"
SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer"
TEXTURE_MATERIAL_NODE_GROUP_NAME = "O Texture Layer"
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Surface Layer"
BALLOON_MATERIAL_NODE_GROUP_NAME = "O Balloon Layer"
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer"
MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer"
SHADELESS_NODE_GROUP_NAME = "O Shadeless"
@@ -250,7 +251,11 @@ def material_node_group(depth_plane=False, surface_group_name=None):
("Alpha", "NodeSocketFloat"),
("Bump Scale", "NodeSocketFloat"),
]
if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME:
normal_group = group_name in {
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME,
BALLOON_MATERIAL_NODE_GROUP_NAME,
}
if normal_group:
input_sockets.insert(2, ("Normal", "NodeSocketColor"))
input_sockets.insert(3, ("Normal Strength", "NodeSocketFloat"))
if depth_plane:
@@ -271,7 +276,9 @@ def material_node_group(depth_plane=False, surface_group_name=None):
socket.min_value = -10000.0
socket.max_value = 10000.0
elif name == "Normal Strength":
socket.default_value = 1.0
socket.default_value = (
0.5 if group_name == BALLOON_MATERIAL_NODE_GROUP_NAME else 1.0
)
socket.min_value = 0.0
socket.max_value = 10.0
output_sockets = [
@@ -304,9 +311,13 @@ def material_node_group(depth_plane=False, surface_group_name=None):
roughness_curve.location = (-250.0, 350.0)
bump = nodes.new("ShaderNodeBump")
bump.location = (-250.0, 0.0)
if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME:
if normal_group:
normal_map = nodes.new("ShaderNodeNormalMap")
normal_map.space = "OBJECT"
normal_map.space = (
"TANGENT"
if group_name == BALLOON_MATERIAL_NODE_GROUP_NAME
else "OBJECT"
)
normal_map.location = (-500.0, -150.0)
links.new(group_input.outputs["Normal"], normal_map.inputs["Color"])
links.new(

46
src/snapform/properties.py
View File

@@ -93,10 +93,7 @@ def models_status(refresh=False):
def _model_files_ready(model):
directory = runtime.storage_root() / "models" / model.directory_name
return all(
any(directory.glob(pattern))
for pattern in model.ready_patterns
)
return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)
def model_ready(model_key):
@@ -147,6 +144,15 @@ def moge2_model_enum_items(_owner, _context):
)
def balloon_depth_model_enum_items(_owner, _context):
models = [
model
for model in model_catalog()
if model.get("family") == "moge2" and model.get("enabled")
]
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
return _model_items(models)
def depth_plane_model_enum_items(_owner, _context):
models = [
model
@@ -165,8 +171,7 @@ def surface_model_enum_items(_owner, _context):
models = [
model
for model in model_catalog()
if model.get("family") in {"moge2", "da3"}
and model.get("enabled")
if model.get("family") in {"moge2", "da3"} and model.get("enabled")
]
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
return _model_items(models)
@@ -228,17 +233,33 @@ def balloon_reuse_material_property():
)
def balloon_ai_refine_property():
def balloon_contour_analysis_property():
return bpy.props.BoolProperty(
name="AI Refine",
name="Contour",
description=(
"Use BEN2 Alpha above 0.5 to refine the lasso before building "
"the Balloon"
"Use BEN2 Alpha above 0.5 to analyze the lasso contour before "
"building the Balloon"
),
default=False,
)
def balloon_depth_analysis_property():
return bpy.props.BoolProperty(
name="Depth",
description=("Use MoGe-2 to add adjustable image depth to the Balloon mesh"),
default=False,
)
def balloon_depth_model_property():
return bpy.props.EnumProperty(
name="Depth Model",
description="MoGe-2 model used for Balloon depth",
items=balloon_depth_model_enum_items,
default=0,
)
def balloon_radial_zoom_property(scale):
return bpy.props.FloatVectorProperty(
size=2,
@@ -247,12 +268,15 @@ def balloon_radial_zoom_property(scale):
options={"HIDDEN", "SKIP_SAVE"},
)
class SnapFormSettings(bpy.types.PropertyGroup):
balloon_mode: balloon_mode_property()
balloon_thickness: balloon_thickness_property()
balloon_resolution: balloon_resolution_property()
balloon_reuse_material: balloon_reuse_material_property()
balloon_ai_refine: balloon_ai_refine_property()
balloon_contour_analysis: balloon_contour_analysis_property()
balloon_depth_analysis: balloon_depth_analysis_property()
balloon_depth_model: balloon_depth_model_property()
balloon_thickness_radial_zoom: balloon_radial_zoom_property(100.0)
balloon_resolution_radial_zoom: balloon_radial_zoom_property(1000.0)
input_path: bpy.props.StringProperty(

8
src/snapform/server/__init__.py
View File

@@ -5,7 +5,7 @@ from blendjob import JobServer
from .jobs import generate_depth as _generate_depth
from .jobs import generate_surface as _generate_surface
from .jobs.background import run as _remove_background
from .jobs.balloon import run as _refine_balloon
from .jobs.balloon import run as _analyze_balloon
from .jobs.debug_ben2 import run as _debug_ben2
from .jobs.debug_da3 import run as _debug_da3
from .jobs.debug_moge2 import run as _debug_moge2
@@ -33,9 +33,9 @@ def download_required_models(context, parameters):
return manager.download_missing(context, parameters["models"])
@server.job("refine-balloon")
def refine_balloon(context, parameters):
return _refine_balloon(context, parameters)
@server.job("analyze-balloon")
def analyze_balloon(context, parameters):
return _analyze_balloon(context, parameters)
@server.job("remove-background")

267
src/snapform/server/jobs/balloon.py
View File

@@ -3,6 +3,7 @@ import math
import time
from pathlib import Path
from ..models import moge2
from ..models.ben2 import infer_foreground
@@ -43,7 +44,7 @@ def create_lasso_input(input_path, polygon, output_path):
masked = Image.fromarray(pixels, mode="RGBA")
output_path = Path(output_path)
masked.save(output_path, format="PNG", compress_level=1)
return output_path, selection_mask, (left, top)
return output_path, selection_mask, (left, top), (width, height)
def largest_component(mask):
@@ -178,15 +179,165 @@ def contour_from_alpha(alpha, selection_mask, origin=(0, 0)):
return [[x + offset_x, y + offset_y] for x, y in contour]
def resized_selection_mask(selection_mask, shape):
import numpy as np
from PIL import Image
selection_mask = np.asarray(selection_mask, dtype=bool)
if selection_mask.shape == tuple(shape):
return selection_mask
resized = Image.fromarray(selection_mask.astype(np.uint8) * 255, mode="L")
resized = resized.resize(
(int(shape[1]), int(shape[0])),
Image.Resampling.NEAREST,
)
return np.asarray(resized, dtype=np.uint8) > 0
def normalized_moge_depth(depth, model_mask, selection_mask):
"""Normalize MoGe depth while preserving its adjustable fitted base plane."""
import numpy as np
depth = np.asarray(depth, dtype=np.float32)
model_mask = np.asarray(model_mask, dtype=bool)
selection_mask = resized_selection_mask(selection_mask, depth.shape)
if model_mask.shape != depth.shape:
raise ValueError("MoGe-2 depth and mask dimensions do not match")
valid = selection_mask & model_mask & np.isfinite(depth) & (depth > 0.0)
if not valid.any():
raise ValueError("MoGe-2 found no valid depth inside the lasso")
height, width = depth.shape
y, x = np.indices(depth.shape, dtype=np.float64)
x = x / max(width - 1, 1) * 2.0 - 1.0
y = y / max(height - 1, 1) * 2.0 - 1.0
design = np.stack((x[valid], y[valid], np.ones(valid.sum())), axis=1)
selected = depth[valid].astype(np.float64)
weights = np.ones_like(selected)
coefficients = np.zeros(3, dtype=np.float64)
for _iteration in range(4):
weighted_design = design * np.sqrt(weights)[:, None]
weighted_depth = selected * np.sqrt(weights)
coefficients = np.linalg.lstsq(
weighted_design,
weighted_depth,
rcond=None,
)[0]
residual = selected - design @ coefficients
centered = residual - np.median(residual)
scale = max(float(np.median(np.abs(centered))) * 1.4826, 1e-6)
weights = np.minimum(
1.0,
(2.5 * scale) / np.maximum(np.abs(centered), 1e-12),
)
baseline = coefficients[0] * x + coefficients[1] * y + coefficients[2]
reference = float(np.median(baseline[valid]))
residual = baseline[valid] - selected
near_residual = np.maximum(residual, 0.0)
near_total = float(np.sum(near_residual))
relative_span = (
float(np.sum(np.square(near_residual))) / near_total
if near_total > 1e-12
else float(np.sqrt(np.mean(np.square(residual))))
)
relative_span = max(relative_span, 1e-6)
depth_signal = np.zeros(depth.shape, dtype=np.float32)
depth_signal[valid] = (reference - selected) / relative_span
baseline_vector = np.asarray(
(
-coefficients[0] / relative_span,
-coefficients[1] / relative_span,
(reference - coefficients[2]) / relative_span,
),
dtype=np.float32,
)
return depth_signal, valid, reference, baseline_vector
def write_depth_analysis(
prediction,
selection_mask,
origin,
source_size,
output_path,
):
import numpy as np
if prediction.depth.shape[0] != 1 or prediction.mask.shape[0] != 1:
raise ValueError("Balloon Depth requires one MoGe-2 frame")
depth_signal, valid, reference, baseline_vector = normalized_moge_depth(
prediction.depth[0],
prediction.mask[0],
selection_mask,
)
crop_size = (selection_mask.shape[1], selection_mask.shape[0])
np.savez_compressed(
output_path,
depth_signal=depth_signal,
baseline_vector=baseline_vector,
valid=valid,
reference_depth=np.asarray(reference, dtype=np.float32),
origin=np.asarray(origin, dtype=np.int32),
crop_size=np.asarray(crop_size, dtype=np.int32),
source_size=np.asarray(source_size, dtype=np.int32),
)
return Path(output_path)
def write_balloon_normal_texture(
normal,
valid,
origin,
crop_size,
source_size,
output_path,
):
"""Write MoGe camera normals as a full-image tangent-space texture."""
import numpy as np
from PIL import Image
normal = np.asarray(normal, dtype=np.float32)
valid = np.asarray(valid, dtype=bool)
if normal.ndim != 3 or normal.shape[2] != 3:
raise ValueError("Balloon normal data must have three components")
if valid.shape != normal.shape[:2]:
valid = resized_selection_mask(valid, normal.shape[:2])
encoded = np.clip(
(normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32) + 1.0)
* 127.5,
0.0,
255.0,
).astype(np.uint8)
encoded[~valid] = (128, 128, 255)
crop_width, crop_height = (int(value) for value in crop_size)
crop = Image.fromarray(encoded, mode="RGB")
if crop.size != (crop_width, crop_height):
crop = crop.resize(
(crop_width, crop_height),
Image.Resampling.BILINEAR,
)
source_width, source_height = (int(value) for value in source_size)
full = Image.new("RGB", (source_width, source_height), (128, 128, 255))
full.paste(crop, tuple(int(value) for value in origin))
output_path = Path(output_path)
full.save(output_path)
return output_path
def run(context, parameters):
"""Refine a Balloon lasso with BEN2 and return its largest contour."""
"""Analyze a Balloon lasso with optional BEN2 contour and MoGe-2 depth."""
total_started = time.perf_counter()
contour_analysis = bool(parameters.get("contour_analysis"))
depth_analysis = bool(parameters.get("depth_analysis"))
if not contour_analysis and not depth_analysis:
raise ValueError("Select Contour or Depth")
context.progress(0.02, "Preparing lasso input")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
input_path, selection_mask, origin = create_lasso_input(
input_path, selection_mask, origin, source_size = create_lasso_input(
parameters["input"],
parameters["polygon"],
output_directory / "lasso-input.png",
@@ -195,41 +346,83 @@ def run(context, parameters):
device = parameters["device"]
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
context.progress(0.18, f"Loading BEN2 for {device.upper()}")
foreground, load_ms, inference_ms = infer_foreground(
input_path,
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
result = {}
if contour_analysis:
model_directory = model_manager.directory(parameters["contour_model"])
context.progress(0.15, f"Loading BEN2 for {device.upper()}")
foreground, _load_ms, _inference_ms = infer_foreground(
input_path,
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(
0.55 if depth_analysis else 0.9,
"Tracing BEN2 foreground",
)
import numpy as np
from PIL import Image
context.progress(0.9, "Tracing BEN2 foreground")
import numpy as np
from PIL import Image
alpha = np.asarray(foreground.getchannel("A"), dtype=np.uint8)
contour = contour_from_alpha(alpha, selection_mask, origin)
gated_alpha = np.where(selection_mask, alpha, 0).astype(np.uint8)
foreground_pixels = foreground.convert("RGBA")
foreground_pixels.putalpha(Image.fromarray(gated_alpha, mode="L"))
foreground_pixels.save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
contour_path = output_directory / "contour.json"
contour_path.write_text(
json.dumps({"points": contour}, separators=(",", ":")),
encoding="utf-8",
)
result["contour"] = "contour.json"
alpha = np.asarray(foreground.getchannel("A"), dtype=np.uint8)
contour = contour_from_alpha(alpha, selection_mask, origin)
gated_alpha = np.where(selection_mask, alpha, 0).astype(np.uint8)
foreground_pixels = foreground.convert("RGBA")
foreground_pixels.putalpha(Image.fromarray(gated_alpha, mode="L"))
foreground_pixels.save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
contour_path = output_directory / "contour.json"
contour_path.write_text(
json.dumps({"points": contour}, separators=(",", ":")),
encoding="utf-8",
)
if depth_analysis:
context.progress(
0.12 if not contour_analysis else 0.58,
"Estimating Balloon depth with MoGe-2",
)
prediction = moge2.infer(
model_manager,
{
"model": parameters["depth_model"],
"device": device,
"resolution_level": parameters.get("resolution_level", 5),
},
[input_path],
include_points=False,
progress_callback=None,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(0.9, "Writing Balloon depth")
depth_path = write_depth_analysis(
prediction,
selection_mask,
origin,
source_size,
output_directory / "depth.npz",
)
result["depth"] = depth_path.name
import numpy as np
with np.load(depth_path, allow_pickle=False) as data:
normal_path = write_balloon_normal_texture(
prediction.normal[0],
data["valid"],
origin,
(selection_mask.shape[1], selection_mask.shape[0]),
source_size,
output_directory / "normal.png",
)
result["normal"] = normal_path.name
total_ms = (time.perf_counter() - total_started) * 1000.0
message = (
f"Balloon contour refined in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
)
message = f"Balloon AI analysis finished in {total_ms:.0f} ms"
context.progress(1.0, message)
return {"contour": "contour.json"}
return result