feat: Add restore cutout depth extension with edge depth correction
- Introduced a new extension for restoring cutout depth with edge correction. - Implemented edge depth extension algorithm in `edge_depth.py` to enhance depth accuracy near cutout edges. - Updated `generate_moge_artifacts` to include an option for extending depth edges during depth generation. - Modified cutout job processing to enable edge depth extension based on depth request. - Added comprehensive tests for edge depth functionality, ensuring correct behavior under various scenarios. - Updated existing node loading functions to remove asset revision dependencies, simplifying the loading process. - Created new design and proposal documentation outlining the goals and decisions for the extension.
This commit is contained in:
@@ -10,12 +10,10 @@ def node_asset_path():
|
||||
return Path(__file__).resolve().parents[1] / "assets" / NODE_ASSET_NAME
|
||||
|
||||
|
||||
def load_node_group(group_name, *, asset_revision=None):
|
||||
def load_node_group(group_name):
|
||||
existing = bpy.data.node_groups.get(group_name)
|
||||
if existing is not None:
|
||||
if asset_revision is None or existing.get("anyimage_asset_revision") == asset_revision:
|
||||
return existing
|
||||
existing.name = group_name + " (Saved)"
|
||||
return existing
|
||||
|
||||
asset_path = node_asset_path()
|
||||
if not asset_path.is_file():
|
||||
@@ -27,7 +25,7 @@ def load_node_group(group_name, *, asset_revision=None):
|
||||
directory=directory,
|
||||
filename=group_name,
|
||||
link=False,
|
||||
do_reuse_local_id=True,
|
||||
do_reuse_local_id=False,
|
||||
)
|
||||
if "FINISHED" not in result:
|
||||
raise RuntimeError("Unable to append the node asset")
|
||||
|
||||
@@ -32,7 +32,7 @@ def create_panorama_from_result(context, result, operator):
|
||||
try:
|
||||
with material_color_image(input_path=operator.color_path) as color:
|
||||
metadata = load_panorama_metadata(result.file("depth_metadata"))
|
||||
group = load_node_group("O Image Depth Panorama", asset_revision=15)
|
||||
group = load_node_group("O Image Depth Panorama")
|
||||
depth = load_depth_result_image(result.file("depth"), source, metadata)
|
||||
material = create_image_material(source.data, color, shadeless=True, scene=context.scene,
|
||||
texture_extension="REPEAT")
|
||||
|
||||
@@ -62,7 +62,7 @@ def create_depth_plane_object(
|
||||
node_group = load_node_group({
|
||||
"DEPTH": "O Image Depth Plane",
|
||||
"RELIEF": "O Image Relief Plane",
|
||||
}[plane_type], asset_revision=16 if plane_type == "DEPTH" else None)
|
||||
}[plane_type])
|
||||
source_name = source_object.name
|
||||
bounds = image_empty_bounds(source_object)
|
||||
mesh = create_image_plane_mesh(source_name, bounds)
|
||||
|
||||
@@ -72,7 +72,7 @@ def _assign_float_point_attribute(mesh, name, values):
|
||||
|
||||
|
||||
def load_cutout_node_group(shape):
|
||||
return load_node_group(CUTOUT_NODE_GROUP_NAMES[shape], asset_revision=26 if shape == "DEPTH_SOLID" else 14 if shape == "DEPTH_BALLOON" else 13)
|
||||
return load_node_group(CUTOUT_NODE_GROUP_NAMES[shape])
|
||||
|
||||
|
||||
def _connect_cutout_normal_scale(material):
|
||||
|
||||
79
src/anyimage/server/geometry/edge_depth.py
Normal file
79
src/anyimage/server/geometry/edge_depth.py
Normal file
@@ -0,0 +1,79 @@
|
||||
"""Extend reliable depth along each visible component's silhouette."""
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
|
||||
def extend_edge_depth(frame, alpha, alpha_threshold=0.9):
|
||||
"""Return edge-corrected geometry while preserving reliable camera points."""
|
||||
import numpy as np
|
||||
from scipy import ndimage
|
||||
|
||||
from ..models.onnx_moge2 import depth_to_points
|
||||
|
||||
depth = frame.depth
|
||||
alpha = np.asarray(alpha, dtype=np.float32)
|
||||
if alpha.shape != depth.shape:
|
||||
raise ValueError("Edge Alpha must match the depth dimensions")
|
||||
if np.all(alpha == 1.0):
|
||||
return frame
|
||||
|
||||
visible = (alpha > 0.0) & (alpha >= alpha_threshold)
|
||||
opaque = visible & (alpha >= 0.95)
|
||||
candidates = (
|
||||
opaque
|
||||
& (frame.validity > 0.5)
|
||||
& np.isfinite(depth)
|
||||
& (depth > 0.0)
|
||||
)
|
||||
if not candidates.any():
|
||||
return frame
|
||||
interior_distance = ndimage.distance_transform_edt(
|
||||
np.pad(opaque, 3, mode="edge")
|
||||
)[3:-3, 3:-3]
|
||||
interior = opaque & (interior_distance > 2.0)
|
||||
core = candidates & interior
|
||||
del interior_distance
|
||||
labels, _ = ndimage.label(visible, structure=np.ones((3, 3), dtype=bool))
|
||||
corrected = depth.copy()
|
||||
background_distance = np.full(depth.shape, np.inf, dtype=np.float32)
|
||||
|
||||
for component_id, bounds in enumerate(ndimage.find_objects(labels), start=1):
|
||||
region = tuple(
|
||||
slice(max(0, bound.start - 7), min(size, bound.stop + 7))
|
||||
for bound, size in zip(bounds, depth.shape)
|
||||
)
|
||||
component = labels[region] == component_id
|
||||
donors = core[region] & component
|
||||
if not donors.any():
|
||||
donors = candidates[region] & component
|
||||
if not donors.any():
|
||||
continue
|
||||
|
||||
indices = ndimage.distance_transform_edt(
|
||||
~donors, return_distances=False, return_indices=True
|
||||
)
|
||||
extended = depth[region][tuple(indices)]
|
||||
smoothed = ndimage.gaussian_filter(extended, sigma=0.8, radius=3)
|
||||
distance = ndimage.distance_transform_edt(~component)
|
||||
background = (
|
||||
~visible[region]
|
||||
& (distance <= 4.0)
|
||||
& (distance < background_distance[region])
|
||||
)
|
||||
target = (component & ~interior[region] & ~donors) | background
|
||||
corrected[region][target] = smoothed[target]
|
||||
background_distance[region][background] = distance[background]
|
||||
|
||||
changed = corrected != depth
|
||||
if not changed.any():
|
||||
return frame
|
||||
points = frame.points
|
||||
if points is not None:
|
||||
height, width = depth.shape
|
||||
intrinsics = frame.intrinsics.copy()
|
||||
intrinsics[0] /= width
|
||||
intrinsics[1] /= height
|
||||
projected = depth_to_points(corrected, intrinsics)
|
||||
points = points.copy()
|
||||
points[changed] = projected[changed]
|
||||
return replace(frame, depth=corrected, points=points)
|
||||
@@ -8,6 +8,7 @@ from .depth_texture import (
|
||||
write_depth_texture,
|
||||
)
|
||||
from .normal_texture import write_camera_normal_texture
|
||||
from .edge_depth import extend_edge_depth
|
||||
|
||||
|
||||
def _limited_model_input(source_path, output_path, max_input_size):
|
||||
@@ -53,6 +54,7 @@ def generate_moge_artifacts(
|
||||
generate_depth=False,
|
||||
normal_mode="NONE",
|
||||
alpha_threshold=0.9,
|
||||
extend_depth_edges=False,
|
||||
):
|
||||
"""Run MoGe once and write only the requested file artifacts."""
|
||||
normal_mode = str(normal_mode).upper()
|
||||
@@ -75,11 +77,15 @@ def generate_moge_artifacts(
|
||||
if generate_depth:
|
||||
context.progress(0.45, "Generating depth texture")
|
||||
alpha = _image_alpha(input_path, frame.depth.shape)
|
||||
depth_frame = (
|
||||
extend_edge_depth(frame, alpha, alpha_threshold)
|
||||
if extend_depth_edges else frame
|
||||
)
|
||||
result["depth"] = write_depth_texture(
|
||||
frame, context.directory / "depth.exr", alpha=alpha,
|
||||
depth_frame, context.directory / "depth.exr", alpha=alpha,
|
||||
)
|
||||
result["depth_metadata"] = write_depth_metadata(
|
||||
frame,
|
||||
depth_frame,
|
||||
(alpha > 0.0) & (alpha >= alpha_threshold),
|
||||
context.directory / "depth.json",
|
||||
)
|
||||
|
||||
@@ -18,6 +18,7 @@ def run(context, parameters):
|
||||
parameters,
|
||||
input_path,
|
||||
generate_depth=generate_depth,
|
||||
extend_depth_edges=generate_depth,
|
||||
normal_mode=normal_mode,
|
||||
alpha_threshold=float(parameters.get("alpha_threshold", 0.9)),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user