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

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)