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

Refactor and enhance DA3 and BEN2 debug panels, introduce mesh export functionality

Browse Source 
- Renamed DA3 and BEN2 debug panels for clarity and consistency.
- Added a new function to draw DA3 input settings to reduce code duplication.
- Implemented mesh export functionality, including confidence thresholding and edge filtering.
- Updated export format options to include new mesh type and adjusted related properties.
- Enhanced job handling to move generated inputs into job root and manage temporary directories.
- Improved error handling and file management during export processes.
- Added unit tests for new mesh functionality and ensured existing tests cover changes.
This commit is contained in:
马南
2026-08-09 21:31:08 +08:00
parent 8aac4822c6
commit 6d4d1c4ec2
10 changed files with 778 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/snapform/__init__.py
View File

@@ -17,8 +17,8 @@ from .operators import (
close_active_job,
)
from .panel import (
Ben2DebugPanel,
Da3DebugPanel,
Ben2RemovalDebugPanel,
Da3OutputDebugPanel,
DebugPanel,
HostPanel,
ModelsPanel,
@@ -42,8 +42,8 @@ CLASSES = (
HostPanel,
ModelsPanel,
DebugPanel,
Da3DebugPanel,
Ben2DebugPanel,
Da3OutputDebugPanel,
Ben2RemovalDebugPanel,
)

46
src/snapform/operators.py
View File

@@ -122,6 +122,7 @@ def close_active_job():
shutil.rmtree(temporary_path, ignore_errors=True)
else:
temporary_path.unlink(missing_ok=True)
temporary_path.parent.rmdir()
except OSError:
pass
_active_job = None
@@ -221,6 +222,12 @@ class RunJob(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
export_format = ""
if self.job_type == RECONSTRUCTION_JOB:
export_format = settings.export_format
elif self.job_type == IMAGE_DEPTH_JOB:
export_format = "mini_npz"
job = ActiveJob(
kind=self.job_type,
process=None,
@@ -228,11 +235,7 @@ class RunJob(bpy.types.Operator):
log_path=log_path,
status_path=paths["status"],
result_directory=paths["result"],
export_format=(
settings.export_format if self.job_type == RECONSTRUCTION_JOB
else "mini_npz" if self.job_type == IMAGE_DEPTH_JOB
else ""
),
export_format=export_format,
source_object_name=self.snapform_source_object_name,
depth_data_name=self.snapform_depth_data_name,
process_res=self.snapform_process_res,
@@ -310,6 +313,21 @@ class RunJob(bpy.types.Operator):
def _prepare_command(self, _context, settings):
paths = runtime.new_job_paths(self.job_type)
if (
getattr(self, "snapform_delete_input", False)
and self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
}
):
archived_input = runtime.move_input_into_job(
self.snapform_input_path,
paths["result"],
)
self.snapform_input_path = str(archived_input)
self.snapform_delete_input = False
background_depth_distance = float(
getattr(self, "snapform_background_depth_distance", 0.0)
)
@@ -389,13 +407,14 @@ class RunJob(bpy.types.Operator):
),
paths,
)
if self.job_type not in {RECONSTRUCTION_JOB, IMAGE_DEPTH_JOB}:
if self.job_type not in {
RECONSTRUCTION_JOB,
IMAGE_DEPTH_JOB,
}:
raise RuntimeError(f"Unsupported task type: {self.job_type}")
export_format = (
"mini_npz"
if self.job_type == IMAGE_DEPTH_JOB
else settings.export_format
)
export_format = settings.export_format
if self.job_type == IMAGE_DEPTH_JOB:
export_format = "mini_npz"
use_ray_pose = None
process_res = None
model_key = selected_model_key(settings, export_format)
@@ -614,7 +633,10 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
if job.kind in {RECONSTRUCTION_JOB, BACKGROUND_REMOVAL_JOB}:
if job.kind in {
RECONSTRUCTION_JOB,
BACKGROUND_REMOVAL_JOB,
}:
try:
open_job_directory(job.result_directory)
except Exception as open_error:

68
src/snapform/panel.py
View File

@@ -160,9 +160,33 @@ class DebugPanel(bpy.types.Panel):
layout.prop(settings, "device")
class Da3DebugPanel(bpy.types.Panel):
bl_label = "DA3 Reconstruction"
bl_idname = "SNAPFORM_PT_debug_da3"
def draw_da3_input_settings(layout, settings, model_property):
layout.prop(settings, model_property)
layout.prop(settings, "input_path")
input_path = bpy.path.abspath(settings.input_path)
if input_path.lower().endswith(
(
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
)
):
layout.prop(settings, "video_frames")
layout.prop(settings, "process_res")
layout.prop(settings, "process_res_method")
layout.prop(settings, "use_ray_pose")
layout.prop(settings, "ref_view_strategy")
class Da3OutputDebugPanel(bpy.types.Panel):
bl_label = "DA3 Output Test"
bl_idname = "SNAPFORM_PT_debug_da3_output"
bl_parent_id = DebugPanel.bl_idname
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
@@ -178,28 +202,7 @@ class Da3DebugPanel(bpy.types.Panel):
model_property = (
"gs_model" if settings.export_format == "gs_ply" else "model"
)
layout.prop(settings, model_property)
layout.prop(settings, "input_path")
input_path = bpy.path.abspath(settings.input_path)
if input_path.lower().endswith(
(
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
)
):
layout.prop(settings, "video_frames")
layout.prop(settings, "process_res")
layout.prop(settings, "process_res_method")
layout.prop(settings, "use_ray_pose")
layout.prop(settings, "ref_view_strategy")
draw_da3_input_settings(layout, settings, model_property)
if settings.export_format == "glb":
glb_box = layout.box()
@@ -211,7 +214,12 @@ class Da3DebugPanel(bpy.types.Panel):
ply_box = layout.box()
ply_box.label(text="PLY (3DGS) Settings")
ply_box.prop(settings, "gs_views_interval")
elif settings.export_format == "mesh":
mesh_box = layout.box()
mesh_box.label(text="Mesh Settings")
mesh_box.prop(settings, "conf_thresh_percentile")
mesh_box.prop(settings, "mesh_filter_edges")
mesh_box.prop(settings, "mesh_max_edge_length")
export_row = layout.row()
export_row.enabled = (
active_job() is None
@@ -220,15 +228,15 @@ class Da3DebugPanel(bpy.types.Panel):
)
operator = export_row.operator(
"snapform.run_job",
text="Export",
text="Generate Output",
icon="EXPORT",
)
operator.job_type = "RECONSTRUCT"
class Ben2DebugPanel(bpy.types.Panel):
bl_label = "BEN2 Background Removal"
bl_idname = "SNAPFORM_PT_debug_ben2"
class Ben2RemovalDebugPanel(bpy.types.Panel):
bl_label = "BEN2 Removal Test"
bl_idname = "SNAPFORM_PT_debug_ben2_removal"
bl_parent_id = DebugPanel.bl_idname
bl_space_type = "VIEW_3D"
bl_region_type = "UI"

48
src/snapform/properties.py
View File

@@ -66,8 +66,8 @@ class SnapFormSettings(bpy.types.PropertyGroup):
default="upper_bound_resize",
)
conf_thresh_percentile: bpy.props.FloatProperty(
name="Conf Thresh Percentile",
description="DA3 GLB confidence threshold percentile",
name="Confidence Threshold",
description="Remove the lowest-confidence percentile of points or mesh faces",
default=40.0,
min=0.0,
max=100.0,
@@ -110,24 +110,52 @@ class SnapFormSettings(bpy.types.PropertyGroup):
default="saddle_balanced",
)
export_format: bpy.props.EnumProperty(
name="Export Format",
description="Official DA3 export format",
name="Output Type",
description="Result generated by the DA3 debug task",
items=(
("glb", "GLB", "Export a colored point cloud as GLB"),
("gs_ply", "PLY (3DGS)", "Export Gaussian splats as PLY"),
(
"glb",
"Point Cloud (GLB)",
"Generate a colored point cloud and optional cameras",
),
(
"gs_ply",
"Gaussian Splats (PLY)",
"Generate a 3D Gaussian Splatting PLY",
),
(
"depth_vis",
"Depth Images",
"Export image and depth visualization pairs",
"Depth Preview (Images)",
"Generate image and colored-depth preview pairs",
),
(
"mini_npz",
"Mini NPZ",
"Export compact depth and camera data",
"Depth Data (NPZ)",
"Generate compact depth, confidence, and camera data",
),
(
"mesh",
"Textured Mesh (GLB)",
"Generate one filtered textured depth mesh per input image",
),
),
default="glb",
)
mesh_filter_edges: bpy.props.BoolProperty(
name="Filter Depth Edges",
description="Remove faces around sharp depth transitions",
default=True,
)
mesh_max_edge_length: bpy.props.FloatProperty(
name="Max Edge Length",
description=(
"Remove faces with an edge longer than this scene distance; "
"zero disables the test"
),
default=0.1,
min=0.0,
soft_max=2.0,
)
gs_views_interval: bpy.props.IntProperty(
name="GS Views Interval",
description="Export 3D Gaussians every N input views",

52
src/snapform/runtime.py
View File

@@ -6,6 +6,7 @@ import shutil
import socket
import subprocess
import sys
import tempfile
import threading
import time
import uuid
@@ -34,9 +35,12 @@ DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = {
"INSTALL_RUNTIME": "install-runtime",
"DOWNLOAD_MODEL": "download-model",
"IMAGE_DEPTH": "depth-plane",
"IMAGE_BACKGROUND": "remove-background",
"ORIGINAL_FRAMES": "depth-plane",
"RECONSTRUCT": "export-reconstruction",
"IMAGE_DEPTH": "generate-depth-plane",
"REMOVE_BACKGROUND": "remove-background",
"IMAGE_BACKGROUND": "remove-image-background",
"ORIGINAL_BACKGROUND": "generate-original-depth-plane",
"ORIGINAL_FRAMES": "extract-original-frames",
}
@@ -196,15 +200,37 @@ def jobs_directory():
def temporary_input_path():
path = jobs_directory() / "inputs"
path.mkdir(parents=True, exist_ok=True)
return path / f"{uuid.uuid4().hex}.png"
directory = Path(tempfile.mkdtemp(prefix="snapform-input-"))
return directory / "input.png"
def temporary_input_directory():
path = jobs_directory() / "inputs" / f"{uuid.uuid4().hex}"
path.mkdir(parents=True, exist_ok=True)
return path
return Path(tempfile.mkdtemp(prefix="snapform-input-"))
def move_input_into_job(input_path, job_directory):
"""Move a Blender-generated input copy into the flat job directory."""
input_path = Path(input_path)
job_directory = Path(job_directory)
if input_path.is_file():
destination = job_directory / "input.png"
input_path.replace(destination)
try:
input_path.parent.rmdir()
except OSError:
pass
return destination
if not input_path.is_dir():
raise ValueError(f"Temporary input does not exist: {input_path}")
for source in input_path.iterdir():
if not source.is_file():
raise ValueError(f"Temporary input contains a directory: {source}")
destination = job_directory / source.name
if destination.exists():
raise ValueError(f"Job input filename already exists: {source.name}")
source.replace(destination)
input_path.rmdir()
return job_directory
def job_folder_name(job_type):
@@ -213,7 +239,7 @@ def job_folder_name(job_type):
def new_job_paths(job_type):
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
stem = f"{job_folder_name(job_type)}_{timestamp}"
directory = jobs_directory() / stem
directory.mkdir(parents=True, exist_ok=True)
@@ -824,6 +850,12 @@ def reconstruct_request(
"ref_view_strategy": settings.ref_view_strategy,
"remove_background_model_dir": "",
"background_depth_distance": float(background_depth_distance),
"mesh_filter_edges": bool(
getattr(settings, "mesh_filter_edges", True)
),
"mesh_max_edge_length": float(
getattr(settings, "mesh_max_edge_length", 0.1)
),
}
if remove_background or background_depth_distance > 0.0:
parameters["remove_background_model_dir"] = str(

320
src/snapform/worker.py
View File

@@ -24,7 +24,7 @@ VIDEO_SUFFIXES = {
".mpeg",
}
SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz", "mesh")
HF_ENDPOINT = "https://huggingface.co"
MODELSCOPE_ENDPOINT = "https://modelscope.cn"
MODELSCOPE_REVISION = "master"
@@ -340,10 +340,211 @@ def copy_frames(images, output_directory):
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
shutil.copy2(
image_path,
Path(output_directory) / f"frame.{index + 1:04d}{suffix}",
destination = (
Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)
def unproject_depth_frame(depth, intrinsics, extrinsics):
import numpy as np
height, width = depth.shape
x, y = np.meshgrid(
np.arange(width, dtype=np.float32),
np.arange(height, dtype=np.float32),
)
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
camera_points = rays * depth.reshape(-1, 1)
camera_points = np.column_stack(
(camera_points, np.ones(camera_points.shape[0], dtype=np.float32))
)
world_to_camera = np.eye(4, dtype=np.float32)
world_to_camera[: extrinsics.shape[0], :4] = extrinsics
world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T
return world_points[:, :3].reshape(height, width, 3)
def depth_edge_mask(depth, threshold=12.0 / 255.0):
import numpy as np
padded = np.pad(depth.astype(np.float32), 1, mode="edge")
grad_x = (
-padded[:-2, :-2]
+ padded[:-2, 2:]
- 2.0 * padded[1:-1, :-2]
+ 2.0 * padded[1:-1, 2:]
- padded[2:, :-2]
+ padded[2:, 2:]
)
grad_y = (
-padded[:-2, :-2]
- 2.0 * padded[:-2, 1:-1]
- padded[:-2, 2:]
+ padded[2:, :-2]
+ 2.0 * padded[2:, 1:-1]
+ padded[2:, 2:]
)
magnitude = np.hypot(grad_x, grad_y)
finite = np.isfinite(magnitude)
if not finite.any():
return np.ones(depth.shape, dtype=bool)
minimum = float(magnitude[finite].min())
maximum = float(magnitude[finite].max())
if maximum <= minimum:
return np.zeros(depth.shape, dtype=bool)
normalized = (magnitude - minimum) / (maximum - minimum)
return normalized >= threshold
def confidence_percentile_threshold(depth, confidence, percentile):
import numpy as np
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
if not valid.any():
raise ValueError("DA3 returned no valid depth confidence values")
return float(np.percentile(confidence[valid], percentile))
def mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold,
filter_edges,
max_edge_length,
):
import numpy as np
height, width = depth.shape
indices = np.arange(height * width, dtype=np.int32).reshape(height, width)
quads = np.stack(
(
indices[:-1, :-1],
indices[1:, :-1],
indices[1:, 1:],
indices[:-1, 1:],
),
axis=-1,
).reshape(-1, 4)
valid = (
np.isfinite(points).all(axis=-1)
& np.isfinite(depth)
& (depth > 0.0)
& (confidence >= confidence_threshold)
)
if filter_edges:
valid &= ~depth_edge_mask(depth)
face_valid = valid.reshape(-1)[quads].all(axis=1)
if max_edge_length > 0.0:
flat_points = points.reshape(-1, 3)
quad_points = flat_points[quads]
next_points = np.roll(quad_points, -1, axis=1)
edge_lengths = np.linalg.norm(quad_points - next_points, axis=-1)
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
quads = quads[face_valid]
if not len(quads):
return np.empty((0, 3), dtype=np.int32)
return np.concatenate(
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
axis=0,
)
def export_textured_mesh_glb(
prediction,
image_paths,
output_directory,
confidence_percentile,
filter_edges,
max_edge_length,
):
import numpy as np
import trimesh
from PIL import Image
required = ("depth", "conf", "intrinsics", "extrinsics")
missing = [name for name in required if getattr(prediction, name) is None]
if missing:
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
threshold = confidence_percentile_threshold(
depth,
confidence,
confidence_percentile,
)
meshes = []
referenced_points = []
for index in range(depth.shape[0]):
points = unproject_depth_frame(
depth[index],
prediction.intrinsics[index],
prediction.extrinsics[index],
)
faces = mesh_triangle_faces(
points,
depth[index],
confidence[index],
threshold,
filter_edges,
max_edge_length,
)
if not len(faces):
continue
height, width = depth[index].shape
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
uu, vv = np.meshgrid(u, v)
uv = np.stack((uu, vv), axis=-1).reshape(-1, 2)
with Image.open(image_paths[index]) as source_image:
texture_image = source_image.convert("RGB").copy()
visual = trimesh.visual.texture.TextureVisuals(
uv=uv,
image=texture_image,
)
mesh = trimesh.Trimesh(
vertices=points.reshape(-1, 3),
faces=faces,
visual=visual,
process=False,
)
mesh.remove_unreferenced_vertices()
meshes.append(mesh)
referenced_points.append(mesh.vertices)
if not meshes:
raise ValueError(
"Mesh filters removed every face; lower Confidence Threshold or disable filters"
)
world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64)
if world_to_camera.shape == (3, 4):
homogeneous = np.eye(4, dtype=np.float64)
homogeneous[:3, :4] = world_to_camera
world_to_camera = homogeneous
cv_to_gltf = np.eye(4, dtype=np.float64)
cv_to_gltf[1, 1] = -1.0
cv_to_gltf[2, 2] = -1.0
alignment = cv_to_gltf @ world_to_camera
aligned_points = trimesh.transform_points(
np.concatenate(referenced_points, axis=0),
alignment,
)
center = np.median(aligned_points, axis=0)
alignment[:3, 3] -= center
scene = trimesh.Scene()
for index, mesh in enumerate(meshes):
mesh.apply_transform(alignment)
scene.add_geometry(mesh, node_name=f"DepthMesh.{index + 1:04d}")
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
output_path = output_directory / "mesh.glb"
scene.export(output_path)
return output_path
class DownloadProgressReporter:
@@ -856,7 +1057,7 @@ def remove_background(args, model_cache=None, cancel_check=None):
images, temporary_directory = prepare_background_input(
args.input,
video_frames,
output_directory.parent,
output_directory,
start_frame=start_frame,
cancel_check=cancel_check,
)
@@ -933,7 +1134,7 @@ def extract_frames(args, cancel_check=None):
images, temporary_directory = prepare_input(
args.input,
getattr(args, "video_frames", 16),
output_directory.parent,
output_directory,
start_frame=getattr(args, "start_frame", 1),
cancel_check=cancel_check,
)
@@ -956,6 +1157,9 @@ def exported_paths(output_directory, export_format):
return sorted((output_directory / "gs_ply").glob("*.ply"))
if export_format == "depth_vis":
return sorted((output_directory / "depth_vis").glob("*.jpg"))
if export_format == "mesh":
result_path = output_directory / "mesh.glb"
return [result_path] if result_path.is_file() else []
if export_format == "mini_npz":
result_path = output_directory / "results.npz"
if result_path.is_file():
@@ -1002,6 +1206,70 @@ def wait_for_exported_paths(
return []
def move_exported_file(
source,
destination,
timeout=10.0,
cancel_check=None,
):
"""Move an async export after Windows releases its file handle."""
source = Path(source)
destination = Path(destination)
deadline = time.monotonic() + timeout
while True:
check_cancelled(cancel_check)
try:
source.replace(destination)
return destination
except OSError as error:
is_file_lock = isinstance(error, PermissionError) or getattr(
error,
"winerror",
None,
) in {32, 33}
if not is_file_lock or time.monotonic() >= deadline:
raise
time.sleep(0.1)
def flatten_exported_paths(
output_directory,
result_paths,
cancel_check=None,
):
"""Move files produced in exporter subdirectories to the job root."""
output_directory = Path(output_directory)
flattened_paths = []
for source in result_paths:
source = Path(source)
check_cancelled(cancel_check)
if source.parent == output_directory:
flattened_paths.append(source)
continue
destination = output_directory / source.name
if destination.exists():
raise RuntimeError(f"Export filename already exists: {destination.name}")
flattened_paths.append(
move_exported_file(
source,
destination,
cancel_check=cancel_check,
)
)
for directory in sorted(
{path.parent for path in map(Path, result_paths)},
key=lambda path: len(path.parts),
reverse=True,
):
while directory != output_directory:
try:
directory.rmdir()
except OSError:
break
directory = directory.parent
return flattened_paths
def reconstruct(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting AI runtime")
check_cancelled(cancel_check)
@@ -1019,7 +1287,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
images, temporary_directory = prepare_input(
args.input,
args.video_frames,
output_directory.parent,
output_directory,
start_frame=getattr(args, "start_frame", 1),
cancel_check=cancel_check,
)
@@ -1124,12 +1392,22 @@ def reconstruct(args, model_cache=None, cancel_check=None):
background_depth_distance,
)
write_status(args.status, 0.88, f"Exporting {args.export_format}")
export(
prediction,
args.export_format,
str(output_directory),
**export_kwargs,
)
if args.export_format == "mesh":
export_textured_mesh_glb(
prediction,
images,
output_directory,
getattr(args, "conf_thresh_percentile", 40.0),
bool(getattr(args, "mesh_filter_edges", True)),
float(getattr(args, "mesh_max_edge_length", 0.1)),
)
else:
export(
prediction,
args.export_format,
str(output_directory),
**export_kwargs,
)
check_cancelled(cancel_check)
result_paths = wait_for_exported_paths(
@@ -1141,19 +1419,13 @@ def reconstruct(args, model_cache=None, cancel_check=None):
raise RuntimeError(
f"No {args.export_format} files were exported"
)
result_paths = flatten_exported_paths(
output_directory,
result_paths,
cancel_check=cancel_check,
)
if args.export_format == "mini_npz":
nested_path = (
output_directory / "exports" / "mini_npz" / "results.npz"
)
result_path = output_directory / "results.npz"
if nested_path.is_file():
nested_path.replace(result_path)
result_paths = [result_path]
for directory in (nested_path.parent, nested_path.parent.parent):
try:
directory.rmdir()
except OSError:
pass
if reference_depth is not None:
write_npz_reference_depth(result_path, reference_depth)
check_cancelled(cancel_check)