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

Add tests for foreground preservation and refactor geometry node asset

Browse Source 
- Implement test to ensure foreground preserves source RGB and Ben2 alpha in WorkerTest.
- Remove outdated tests related to background distance and original background depth.
- Update geometry node asset to reflect new naming conventions and structure.
- Introduce validation for depth plane asset, ensuring correct input/output configurations and attribute handling.
- Adjust confidence and edge parameters in geometry node asset to support a wider range of values.
This commit is contained in:
马南
2026-08-11 23:25:51 +08:00
parent 07d836d9ad
commit 7cd5681a8b
15 changed files with 1296 additions and 660 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
src/snapform/worker.py
View File

@@ -48,64 +48,6 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def apply_foreground_depth_distance(
prediction,
foreground_paths,
background_depth_distance,
):
"""Set fully transparent pixels to a fixed negative displacement."""
import numpy as np
from PIL import Image
if len(foreground_paths) != len(prediction.depth):
raise RuntimeError("Foreground mask count does not match depth frames")
masks = []
height, width = prediction.depth.shape[-2:]
for foreground_path in foreground_paths:
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize((width, height), resample=Image.Resampling.NEAREST)
masks.append(np.asarray(alpha) > 0)
apply_depth_background_array(
prediction,
np.stack(masks),
background_depth_distance,
)
def apply_depth_background_array(
prediction,
foreground_mask,
background_depth_distance,
):
import numpy as np
depth = np.asarray(prediction.depth)
masked_depth = depth.copy()
background_depth = -background_depth_distance
for frame_index in range(len(masked_depth)):
np.maximum(
masked_depth[frame_index],
background_depth,
out=masked_depth[frame_index],
)
masked_depth[frame_index][~foreground_mask[frame_index]] = (
background_depth
)
set_depth_border(masked_depth[frame_index], background_depth)
prediction.depth = masked_depth
def set_depth_border(depth, background_depth):
"""Set the outermost pixel ring to the fixed background depth."""
if depth.ndim != 2 or not depth.size:
raise ValueError("Depth frame must be a non-empty 2D array")
depth[0, :] = background_depth
depth[-1, :] = background_depth
depth[:, 0] = background_depth
depth[:, -1] = background_depth
def convert_ai_depth_for_displacement(prediction):
"""Convert DA3 distance into Blender's positive displacement direction."""
import numpy as np
@@ -137,55 +79,6 @@ def write_npz_reference_depth(npz_path, reference_depth):
temporary_path.replace(npz_path)
def write_original_background_depth(
image_paths,
foreground_paths,
background_depth_distance,
output_directory,
):
"""Write original-image luminance with a fixed negative background depth."""
import numpy as np
from PIL import Image
if len(image_paths) != len(foreground_paths):
raise RuntimeError("Foreground mask count does not match source frames")
depth_frames = []
for image_path, foreground_path in zip(image_paths, foreground_paths):
with Image.open(image_path) as source:
luminance = np.asarray(source.convert("L"), dtype=np.float32) / 255.0
with Image.open(foreground_path) as foreground:
alpha = np.asarray(
foreground.convert("RGBA").getchannel("A"),
dtype=np.uint8,
)
if luminance.shape != alpha.shape:
raise RuntimeError("Foreground mask size does not match source frame")
depth_frames.append(
compose_original_background_depth(
luminance,
alpha,
background_depth_distance,
)
)
np.savez_compressed(
Path(output_directory) / "original-depth.npz",
depth=np.stack(depth_frames),
)
def compose_original_background_depth(
luminance,
alpha,
background_depth_distance,
):
import numpy as np
background_depth = -background_depth_distance
depth = np.where(alpha > 0, luminance, background_depth)
set_depth_border(depth, background_depth)
return depth
def write_status(
path,
progress,
@@ -608,7 +501,7 @@ def export_image_mesh_data(
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")
raise ValueError("Convert to Surface currently supports one image")
raw_depth = depth_frames[0]
confidence = confidence_frames[0]
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
@@ -636,7 +529,7 @@ def export_image_mesh_data(
)
if not len(faces):
raise ValueError(
"DA3 output does not contain any valid mesh faces"
"DA3 output does not contain any valid surface faces"
)
progress(5, "Projecting UV and structured attributes")
vertices, faces = compact_mesh(
@@ -658,7 +551,7 @@ def export_image_mesh_data(
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")
progress(6, "Writing surface data")
result_path = output_directory / "mesh.npz"
np.savez_compressed(
result_path,
@@ -876,7 +769,7 @@ def export_textured_mesh_glb(
referenced_points.append(mesh.vertices)
if not meshes:
raise ValueError(
"Mesh filters removed every face; lower Confidence Threshold or disable filters"
"Surface filters removed every face; lower Confidence Threshold or disable filters"
)
world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64)
@@ -1343,6 +1236,7 @@ def infer_foreground(
image,
refine_foreground=False,
)
foreground = preserve_source_rgb(image, foreground)
if device.type == "cuda":
torch.cuda.synchronize()
inference_ms = (time.perf_counter() - inference_started) * 1000.0
@@ -1354,6 +1248,18 @@ def infer_foreground(
return foreground, load_ms, inference_ms
def preserve_source_rgb(source, foreground):
"""Combine the untouched source RGB with BEN2's foreground alpha."""
source_rgba = source.convert("RGBA")
alpha = foreground.convert("RGBA").getchannel("A")
if alpha.size != source_rgba.size:
from PIL import Image
alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS)
source_rgba.putalpha(alpha)
return source_rgba
def infer_foreground_frames(
images,
model_dir,
@@ -1456,21 +1362,6 @@ def remove_background(args, model_cache=None, cancel_check=None):
cancel_check=cancel_check,
)
copy_frames(images, output_directory)
background_depth_distance = float(
getattr(args, "background_depth_distance", 0.0)
)
if background_depth_distance > 0.0:
foreground_paths = (
[output_directory / "foreground.png"]
if len(images) == 1
else sorted(output_directory.glob("foreground.*.png"))
)
write_original_background_depth(
images,
foreground_paths,
background_depth_distance,
output_directory,
)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -1662,9 +1553,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
"remove_background_model_dir",
"",
)
background_depth_distance = float(
getattr(args, "background_depth_distance", 0.0)
)
foreground_images = []
if remove_background_model_dir:
write_status(args.status, 0.12, "Removing image background")
@@ -1747,12 +1635,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
if args.export_format == "mini_npz":
reference_depth = depth_reference_value(prediction.depth[0])
convert_ai_depth_for_displacement(prediction)
if background_depth_distance > 0.0:
apply_foreground_depth_distance(
prediction,
foreground_images,
background_depth_distance,
)
write_status(args.status, 0.88, f"Exporting {args.export_format}")
if args.export_format == "mesh":
export_textured_mesh_glb(