nan
a0dd6e4a88
- 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.
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""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)
|