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

feat: Add surface generation and media handling

Browse Source 
- Implemented surface generation functionality in `surface.py`, including geometry estimation and texture writing.
- Introduced media input handling in `input.py`, supporting image and video frame extraction.
- Added model management for geometry adapters in `models/__init__.py`.
- Created BEN2 and DA3 model inference logic in `ben2.py` and `da3.py`, respectively.
- Developed ONNX model inference utilities for BEN2 and DA3 in `onnx_ben2.py` and `onnx_da3.py`.
- Added support for MoGe-2 model inference in `moge2.py` with necessary ONNX utilities.
- Implemented output handling for depth and surface data in `depth.py` and `surface.py`.
- Created debug output utilities for DA3 and MoGe-2 predictions in `debug_da3.py` and `debug_moge2.py`.
- Enhanced ONNX runtime session management in `onnx_runtime.py`.
- Updated `__init__.py` files to include new modules and functions for better organization.
This commit is contained in:
马南
2026-08-18 23:27:10 +08:00
parent 9dfa49602c
commit 835225d264
62 changed files with 2126 additions and 2318 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/snapform/server/__init__.py
View File

@@ -2,11 +2,14 @@
from blendjob import JobServer
from .jobs import generate_depth as _generate_depth
from .jobs import generate_surface as _generate_surface
from .jobs.background import run as _remove_background
from .jobs.debug_ben2 import run as _debug_ben2
from .jobs.debug_da3 import run as _debug_da3
from .jobs.debug_moge2 import run as _debug_moge2
from .jobs.frames import run as _extract_frames
from .model_manager import ModelManager, ModelSessionCache
from .background import remove_background as _remove_background
from .da3 import generate_da3 as _generate_da3
from .media import extract_frames as _extract_frames
from .moge2 import generate_moge2 as _generate_moge2
server = JobServer("SnapForm Job Server")
@@ -29,7 +32,6 @@ def download_required_models(context, parameters):
return manager.download_missing(context, parameters["models"])
@server.job("debug-ben2")
@server.job("remove-background")
def remove_background(context, parameters):
return _remove_background(context, parameters)
@@ -40,18 +42,29 @@ def extract_frames(context, parameters):
return _extract_frames(context, parameters)
@server.job("debug-moge2")
@server.job("generate-moge2-surface")
@server.job("generate-moge2-depth")
def generate_moge2(context, parameters):
return _generate_moge2(context, parameters)
@server.job("generate-depth")
def generate_depth(context, parameters):
return _generate_depth(context, parameters)
@server.job("generate-surface")
def generate_surface(context, parameters):
return _generate_surface(context, parameters)
@server.job("debug-da3")
@server.job("generate-da3-surface")
@server.job("generate-da3-depth")
def generate_da3(context, parameters):
return _generate_da3(context, parameters)
def debug_da3(context, parameters):
return _debug_da3(context, parameters)
@server.job("debug-moge2")
def debug_moge2(context, parameters):
return _debug_moge2(context, parameters)
@server.job("debug-ben2")
def debug_ben2(context, parameters):
return _debug_ben2(context, parameters)
__all__ = ("ModelManager", "ModelSessionCache", "server")

170
src/snapform/server/background.py
View File

@@ -1,170 +0,0 @@
import shutil
import time
from pathlib import Path
from .media import (
IMAGE_SUFFIXES,
check_cancelled,
copy_frames,
prepare_background_input,
)
def infer_foreground(
input_path,
model_dir,
device,
model_cache=None,
cancel_check=None,
):
from PIL import Image
input_path = Path(input_path)
check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
from .onnx_ben2 import create_session, infer
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
if owns_model:
load_started = time.perf_counter()
model = create_session(model_dir, requested_device)
load_ms = (time.perf_counter() - load_started) * 1000.0
else:
model, load_ms = model_cache.get_ben2(model_dir, device)
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
inference_started = time.perf_counter()
foreground = infer(model, image)
foreground = preserve_source_rgb(image, foreground)
inference_ms = (time.perf_counter() - inference_started) * 1000.0
check_cancelled(cancel_check)
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,
device,
output_directory,
progress_callback,
start_progress,
progress_span,
model_cache=None,
cancel_check=None,
):
"""Run BEN2 on every frame and write foreground.0001.png etc."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
progress_callback(
start_progress + progress_span * index / len(images),
f"Removing background ({index + 1}/{len(images)})",
)
foreground, load_ms, inference_ms = infer_foreground(
image_path,
model_dir,
device,
model_cache=model_cache,
cancel_check=cancel_check,
)
output_path = output_directory / f"foreground.{index + 1:04d}.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError(f"BEN2 did not write: {output_path}")
total_load_ms += load_ms
total_inference_ms += inference_ms
return total_load_ms, total_inference_ms
def remove_background(context, parameters):
"""Run BEN2 and write transparent foreground images into the Job directory."""
total_started = time.perf_counter()
context.progress(0.02, "Starting Environment")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
device = parameters["device"]
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
video_frames = parameters.get("video_frames", 16)
start_frame = parameters.get("start_frame", 1)
images, temporary_directory = prepare_background_input(
parameters["input"],
video_frames,
output_directory,
start_frame=start_frame,
cancel_check=context.check_cancelled,
)
try:
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
context.progress(0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
else:
load_ms, inference_ms = infer_foreground_frames(
images,
model_directory,
device,
output_directory,
context.progress,
0.12,
0.7,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
context.check_cancelled()
total_ms = (time.perf_counter() - total_started) * 1000.0
frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
message = (
f"Background removed{frame_text} in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
)
print(message)
context.progress(1.0, message)
if tuple(output_directory.glob("foreground.*.png")):
return {"foreground": "."}
return {"foreground": "foreground.png"}

232
src/snapform/server/da3.py
View File

@@ -1,232 +0,0 @@
import shutil
from .background import infer_foreground, infer_foreground_frames
from .geometry import (
export_da3_depth_preview,
export_da3_mini_npz,
export_da3_point_cloud,
export_image_mesh_data,
export_prediction_attributes,
export_textured_mesh_glb,
)
from .media import (
copy_frames,
flatten_exported_paths,
prepare_input,
wait_for_exported_paths,
)
def merge_da3_frame_predictions(predictions):
"""Combine independent single-frame DA3 results for existing exporters."""
import numpy as np
if not predictions:
raise ValueError("DA3 returned no frame predictions")
merged = predictions[0]
for name in (
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
):
values = [getattr(prediction, name) for prediction in predictions]
setattr(merged, name, np.concatenate(values, axis=0))
return merged
def generate_da3(context, parameters):
"""Run DA3 inference and return the artifact used by the requesting Job."""
context.progress(0.02, "Starting Environment")
context.check_cancelled()
from .onnx_da3 import infer
from .onnx_runtime import provider_label
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
context.progress( 0.08, "Preparing input")
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
device = parameters["device"]
export_format = parameters["export_format"]
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
output_directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
context.check_cancelled()
remove_background_model = parameters.get("remove_background_model", "")
remove_background_model_dir = (
model_manager.directory(remove_background_model)
if remove_background_model
else ""
)
foreground_images = []
if remove_background_model_dir:
context.progress( 0.12, "Removing image background")
if len(images) == 1:
foreground, load_ms, inference_ms = infer_foreground(
images[0],
remove_background_model_dir,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
foreground.convert("RGBA").save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
if not (output_directory / "foreground.png").is_file():
raise RuntimeError("BEN2 did not write a foreground image")
foreground_images = [output_directory / "foreground.png"]
else:
load_ms, inference_ms = infer_foreground_frames(
images,
remove_background_model_dir,
device,
output_directory,
context.progress,
0.12,
0.05,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
foreground_images = sorted(
output_directory.glob("foreground.*.png")
)
if remove_background_model_dir:
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)
context.progress( 0.15, "Loading DA3 ONNX model")
model = model_manager.get_da3(model_directory, device)
context.progress(
0.20,
f"Loaded DA3 on {provider_label(model)}",
)
context.check_cancelled()
context.progress(
0.25,
f"Estimating depth for {len(images)} frame(s)",
)
inference_options = {
"process_res": parameters["process_res"],
"process_res_method": parameters["process_res_method"],
"ref_view_strategy": parameters["ref_view_strategy"],
}
try:
if parameters.get("inference_mode", "joint") == "per_frame":
predictions = []
for index, image in enumerate(images):
context.check_cancelled()
context.progress(
0.25 + 0.60 * index / len(images),
f"Estimating DA3 depth ({index + 1}/{len(images)})",
)
predictions.append(
infer(model, [str(image)], **inference_options)
)
prediction = merge_da3_frame_predictions(predictions)
else:
prediction = infer(
model,
[str(path) for path in images],
**inference_options,
)
except UnicodeDecodeError as error:
# DirectML reports failures with a locale-specific message that
# ONNX Runtime cannot always decode, masking the real error.
if provider_label(model) == "DirectML":
raise RuntimeError(
"DA3 inference failed on DirectML. Switch the device to "
"CPU in SnapForm Preferences and retry. "
f"Original error: {error}"
) from error
raise
context.check_cancelled()
context.progress( 0.88, f"Exporting {export_format}")
if export_format == "mesh":
export_textured_mesh_glb(
prediction,
images,
output_directory,
parameters.get("conf_thresh_percentile", 40.0),
bool(parameters.get("mesh_filter_edges", True)),
float(parameters.get("mesh_max_edge_length", 0.1)),
)
elif export_format == "mesh_data":
def mesh_progress(step, total, label):
context.progress(
0.88 + 0.09 * step / total,
f"Preparing surface data {step}/{total}: {label}",
)
export_image_mesh_data(
prediction,
images[0],
output_directory,
foreground_path=(
foreground_images[0] if foreground_images else None
),
progress_callback=mesh_progress,
)
elif export_format == "mini_npz":
export_da3_mini_npz(prediction, output_directory)
elif export_format == "depth_vis":
export_da3_depth_preview(prediction, output_directory)
elif export_format == "glb":
export_da3_point_cloud(
prediction,
output_directory,
parameters["conf_thresh_percentile"],
parameters["num_max_points"],
)
else:
raise ValueError(
f"DA3 ONNX does not support export format: {export_format}"
)
if export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
context.progress(
0.97,
"Writing structured debug attributes",
)
export_prediction_attributes(prediction, output_directory)
context.check_cancelled()
result_paths = wait_for_exported_paths(
output_directory,
export_format,
cancel_check=context.check_cancelled,
)
if not result_paths:
raise RuntimeError(
f"No {export_format} files were exported"
)
result_paths = flatten_exported_paths(
output_directory,
result_paths,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
copy_frames(images, output_directory)
context.progress(
1.0,
f"Exported {len(result_paths)} {export_format} file(s)",
)
if context.job_type == "generate-da3-surface":
return {"surface": "mesh.npz"}
if context.job_type == "generate-da3-depth":
return {"depth": "results.npz"}
return {"directory": "."}
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

616
src/snapform/server/geometry.py
View File

@@ -1,616 +0,0 @@
from pathlib import Path
SCENE_FILENAME = "scene.glb"
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 unproject_camera_depth(depth, intrinsics):
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
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)
def compact_mesh(vertices, faces):
import numpy as np
vertices = np.asarray(vertices, dtype=np.float32)
faces = np.asarray(faces, dtype=np.int32)
used_vertices = np.unique(faces)
index_map = np.full(len(vertices), -1, dtype=np.int32)
index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32)
vertices = vertices[used_vertices]
faces = index_map[faces]
return vertices, faces
def project_camera_uv(vertices, intrinsics, image_size):
import numpy as np
width, height = image_size
z = vertices[:, 2]
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
raise ValueError("Simplified mesh contains invalid camera depth")
u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2]
v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2]
denominator_x = max(width - 1, 1)
denominator_y = max(height - 1, 1)
return np.column_stack(
(
np.clip(u / denominator_x, 0.0, 1.0),
np.clip(1.0 - v / denominator_y, 0.0, 1.0),
)
).astype(np.float32)
def project_camera_pixels(vertices, intrinsics):
import numpy as np
z = vertices[:, 2]
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
raise ValueError("Mesh contains invalid camera depth")
return np.column_stack(
(
intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2],
intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2],
)
)
def sample_image_nearest(values, pixels):
import numpy as np
height, width = values.shape
x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1)
y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1)
return np.asarray(values)[y, x]
def confidence_percentile_map(depth, confidence):
import numpy as np
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
result = np.zeros(depth.shape, dtype=np.float32)
if not valid.any():
return result
sorted_values = np.sort(np.asarray(confidence)[valid])
ranks = np.searchsorted(
sorted_values,
np.asarray(confidence)[valid],
side="right",
)
result[valid] = ranks.astype(np.float32) / len(sorted_values)
return result
def depth_discontinuity_ratio(depth):
import numpy as np
padded = np.pad(depth, 1, mode="edge")
windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3))
difference = np.nanmax(windows, axis=(-2, -1))
difference -= np.nanmin(windows, axis=(-2, -1))
return difference / np.maximum(depth, 1e-6)
def triangle_edge_stretch_values(points, triangles):
import numpy as np
flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3)
triangle_points = flat_points[triangles]
triangle_depth = triangle_points[..., 2]
safe_depth = np.where(
np.abs(triangle_depth) > 1e-8,
triangle_depth,
np.nan,
)
triangle_rays = triangle_points / safe_depth[..., None]
next_points = np.roll(triangle_points, -1, axis=1)
next_rays = np.roll(triangle_rays, -1, axis=1)
next_depth = np.roll(triangle_depth, -1, axis=1)
local_depth = (triangle_depth + next_depth) * 0.5
flat_edges = (
triangle_rays * local_depth[..., None]
- next_rays * local_depth[..., None]
)
flat_lengths = np.linalg.norm(flat_edges, axis=-1)
actual_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
stretch = actual_lengths / np.maximum(flat_lengths, 1e-8)
stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0
return stretch
def image_mesh_attributes(
vertices,
faces,
intrinsics,
confidence_percentiles,
face_depth_discontinuity,
):
import numpy as np
point_percentile = sample_image_nearest(
confidence_percentiles,
project_camera_pixels(vertices, intrinsics),
).astype(np.float32)
face_confidence = point_percentile[faces].min(axis=1)
return {
"point_depth": vertices[:, 2].astype(np.float32),
"face_confidence_percentile": face_confidence.astype(np.float32),
"face_depth_discontinuity": np.asarray(
face_depth_discontinuity,
dtype=np.float32,
),
}
def export_prediction_attributes(prediction, output_directory):
import numpy as np
values = {}
for name in ("depth", "conf", "intrinsics", "extrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.asarray(value)
if "depth" in values and "conf" in values:
values["confidence_percentile"] = np.stack(
[
confidence_percentile_map(depth, confidence)
for depth, confidence in zip(
values["depth"],
values["conf"],
)
]
)
output_path = Path(output_directory) / "prediction.attributes.npz"
np.savez(output_path, **values)
return output_path
def export_da3_mini_npz(prediction, output_directory):
import numpy as np
values = {"depth": np.round(prediction.depth, 8)}
for name in ("conf", "extrinsics", "intrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.round(value, 2) if name == "conf" else value
output_path = Path(output_directory) / "results.npz"
np.savez_compressed(output_path, **values)
return output_path
def export_da3_depth_preview(prediction, output_directory):
import numpy as np
from PIL import Image
target = Path(output_directory) / "depth_vis"
target.mkdir(parents=True, exist_ok=True)
for index, depth in enumerate(prediction.depth):
valid = np.isfinite(depth) & (depth > 0.0)
normalized = np.zeros(depth.shape, dtype=np.float32)
if valid.any():
low, high = np.percentile(depth[valid], (2.0, 98.0))
normalized[valid] = np.clip(
(depth[valid] - low) / max(float(high - low), 1e-6),
0.0,
1.0,
)
color = np.stack(
(
normalized,
1.0 - np.abs(normalized * 2.0 - 1.0),
1.0 - normalized,
),
axis=-1,
)
color = np.asarray(color * 255.0, dtype=np.uint8)
source = np.asarray(prediction.processed_images[index], dtype=np.uint8)
combined = np.concatenate((source, color), axis=1)
Image.fromarray(combined, mode="RGB").save(
target / f"{index:04d}.jpg",
quality=95,
)
def export_da3_point_cloud(
prediction,
output_directory,
confidence_percentile=40.0,
num_max_points=1_000_000,
):
import numpy as np
import trimesh
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
threshold = confidence_percentile_threshold(
depth,
confidence,
confidence_percentile,
)
point_sets = []
color_sets = []
for index in range(depth.shape[0]):
points = unproject_depth_frame(
depth[index],
prediction.intrinsics[index],
prediction.extrinsics[index],
)
valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0)
if threshold is not None:
valid &= confidence[index] >= threshold
point_sets.append(points[valid])
color_sets.append(prediction.processed_images[index][valid])
points = np.concatenate(point_sets, axis=0)
colors = np.concatenate(color_sets, axis=0)
maximum = max(int(num_max_points), 1)
if len(points) > maximum:
indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64)
points = points[indices]
colors = colors[indices]
scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors))
scene.export(Path(output_directory) / SCENE_FILENAME)
def remove_background_faces(faces, foreground_mask):
"""Return only triangles whose vertices all belong to the foreground."""
foreground_points = foreground_mask.reshape(-1)
return faces[foreground_points[faces].all(axis=1)]
def export_image_mesh_data(
prediction,
image_path,
output_directory,
foreground_path=None,
progress_callback=None,
):
import numpy as np
from PIL import Image
total_steps = 6
def progress(step, label):
if progress_callback is not None:
progress_callback(step, total_steps, label)
progress(1, "Validating DA3 data")
required = ("depth", "conf", "intrinsics")
missing = [
name
for name in required
if getattr(prediction, name, None) is None
]
if missing:
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
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 Surface currently supports one image")
raw_depth = depth_frames[0]
confidence = confidence_frames[0]
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
progress(2, "Unprojecting the depth grid")
points = unproject_camera_depth(raw_depth, intrinsics)
progress(3, "Preparing confidence and background data")
foreground_mask = np.ones(raw_depth.shape, dtype=bool)
if foreground_path is not None:
foreground_mask = foreground_mesh_mask(
foreground_path,
raw_depth.shape,
)
confidence_percentiles = confidence_percentile_map(
raw_depth,
confidence,
)
progress(4, "Building source triangles")
faces = mesh_triangle_faces(
points,
raw_depth,
confidence,
None,
False,
0.0,
)
if foreground_path is not None:
faces = remove_background_faces(faces, foreground_mask)
if not len(faces):
raise ValueError(
"DA3 output does not contain any valid surface faces"
)
progress(5, "Preparing source UV and attributes")
depth_discontinuity = depth_discontinuity_ratio(raw_depth)
face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
vertices, faces = compact_mesh(
points.reshape(-1, 3),
faces,
)
height, width = raw_depth.shape
uv = project_camera_uv(vertices, intrinsics, (width, height))
attributes = image_mesh_attributes(
vertices,
faces,
intrinsics,
confidence_percentiles,
face_depth_discontinuity,
)
reference_depth = float(np.median(vertices[:, 2]))
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
progress(6, "Writing source surface data")
result_path = output_directory / "mesh.npz"
np.savez_compressed(
result_path,
vertices=vertices,
faces=faces,
uv=uv,
intrinsics=intrinsics,
image_size=np.asarray((width, height), dtype=np.int32),
reference_depth=np.asarray(reference_depth, dtype=np.float32),
**attributes,
)
texture_path = output_directory / "texture.png"
texture_source = foreground_path or image_path
with Image.open(texture_source) as image:
image.convert("RGBA").save(texture_path, format="PNG", compress_level=1)
return result_path, texture_path
def depth_edge_mask(depth, threshold=12.0 / 255.0):
import numpy as np
values = np.asarray(depth, dtype=np.float32)
edges = np.zeros(values.shape, dtype=bool)
for axis in (0, 1):
first = np.take(values, range(values.shape[axis] - 1), axis=axis)
second = np.take(values, range(1, values.shape[axis]), axis=axis)
valid = (
np.isfinite(first)
& np.isfinite(second)
& (first > 0.0)
& (second > 0.0)
)
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
discontinuity = valid & (np.abs(first - second) / scale > threshold)
lower = [slice(None), slice(None)]
upper = [slice(None), slice(None)]
lower[axis] = slice(0, -1)
upper[axis] = slice(1, None)
edges[tuple(lower)] |= discontinuity
edges[tuple(upper)] |= discontinuity
return edges
def edge_culling_stretch_limit(culling):
factor = min(max(float(culling) / 100.0, 0.0001), 1.0)
return 1.0 / factor
def foreground_mesh_mask(foreground_path, depth_shape):
import numpy as np
from PIL import Image, ImageFilter
height, width = depth_shape
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize(
(width, height),
resample=Image.Resampling.BILINEAR,
)
alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75))
return np.asarray(alpha) >= 128
def confidence_percentile_threshold(depth, confidence, percentile):
import numpy as np
if float(percentile) <= 0.0:
return None
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,
edge_threshold=12.0 / 255.0,
edge_stretch_limit=None,
vertex_mask=None,
):
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)
)
if confidence_threshold is not None:
valid &= np.isfinite(confidence) & (
confidence >= confidence_threshold
)
if filter_edges:
valid &= ~depth_edge_mask(depth, edge_threshold)
if vertex_mask is not None:
if np.shape(vertex_mask) != depth.shape:
raise ValueError("Foreground mask and depth dimensions do not match")
valid &= np.asarray(vertex_mask, dtype=bool)
triangles = np.concatenate(
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
axis=0,
)
face_valid = valid.reshape(-1)[triangles].all(axis=1)
if edge_stretch_limit is not None:
face_valid &= triangle_edge_stretch_mask(
points,
triangles,
edge_stretch_limit,
)
if max_edge_length > 0.0:
flat_points = points.reshape(-1, 3)
triangle_points = flat_points[triangles]
next_points = np.roll(triangle_points, -1, axis=1)
edge_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
triangles = triangles[face_valid]
if not len(triangles):
return np.empty((0, 3), dtype=np.int32)
return triangles
def triangle_edge_stretch_mask(points, triangles, stretch_limit):
stretch = triangle_edge_stretch_values(points, triangles)
tolerance = float(stretch_limit) * (1.0 + 1e-5)
return (stretch <= tolerance).all(axis=1)
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.vertex_attributes["o_depth"] = depth[index].reshape(-1)
mesh.vertex_attributes["o_confidence"] = confidence[
index
].reshape(-1)
mesh.vertex_attributes["o_confidence_percentile"] = (
confidence_percentile_map(
depth[index],
confidence[index],
).reshape(-1)
)
mesh.remove_unreferenced_vertices()
meshes.append(mesh)
referenced_points.append(mesh.vertices)
if not meshes:
raise ValueError(
"Surface 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

35
src/snapform/server/geometry/__init__.py Normal file
View File

@@ -0,0 +1,35 @@
from .grid_mesh import (
compact_mesh,
confidence_percentile_map,
confidence_percentile_threshold,
depth_edge_mask,
depth_discontinuity_ratio,
edge_culling_stretch_limit,
foreground_mesh_mask,
mesh_triangle_faces,
remove_background_faces,
triangle_edge_stretch_values,
triangle_edge_stretch_mask,
)
from .projection import unproject_camera_depth, unproject_depth_frame
from .surface import GeometryFrame, SurfaceMesh, build_surface
__all__ = (
"GeometryFrame",
"SurfaceMesh",
"build_surface",
"compact_mesh",
"confidence_percentile_map",
"confidence_percentile_threshold",
"depth_edge_mask",
"depth_discontinuity_ratio",
"edge_culling_stretch_limit",
"foreground_mesh_mask",
"mesh_triangle_faces",
"remove_background_faces",
"triangle_edge_stretch_values",
"triangle_edge_stretch_mask",
"unproject_camera_depth",
"unproject_depth_frame",
)

202
src/snapform/server/geometry/grid_mesh.py Normal file
View File

@@ -0,0 +1,202 @@
def compact_mesh(vertices, faces):
import numpy as np
vertices = np.asarray(vertices, dtype=np.float32)
faces = np.asarray(faces, dtype=np.int32)
used_vertices = np.unique(faces)
index_map = np.full(len(vertices), -1, dtype=np.int32)
index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32)
vertices = vertices[used_vertices]
faces = index_map[faces]
return vertices, faces
def confidence_percentile_map(depth, confidence):
import numpy as np
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
result = np.zeros(depth.shape, dtype=np.float32)
if not valid.any():
return result
sorted_values = np.sort(np.asarray(confidence)[valid])
ranks = np.searchsorted(
sorted_values,
np.asarray(confidence)[valid],
side="right",
)
result[valid] = ranks.astype(np.float32) / len(sorted_values)
return result
def depth_discontinuity_ratio(depth):
import numpy as np
padded = np.pad(depth, 1, mode="edge")
windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3))
difference = np.nanmax(windows, axis=(-2, -1))
difference -= np.nanmin(windows, axis=(-2, -1))
return difference / np.maximum(depth, 1e-6)
def triangle_edge_stretch_values(points, triangles):
import numpy as np
flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3)
triangle_points = flat_points[triangles]
triangle_depth = triangle_points[..., 2]
safe_depth = np.where(
np.abs(triangle_depth) > 1e-8,
triangle_depth,
np.nan,
)
triangle_rays = triangle_points / safe_depth[..., None]
next_points = np.roll(triangle_points, -1, axis=1)
next_rays = np.roll(triangle_rays, -1, axis=1)
next_depth = np.roll(triangle_depth, -1, axis=1)
local_depth = (triangle_depth + next_depth) * 0.5
flat_edges = (
triangle_rays * local_depth[..., None]
- next_rays * local_depth[..., None]
)
flat_lengths = np.linalg.norm(flat_edges, axis=-1)
actual_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
stretch = actual_lengths / np.maximum(flat_lengths, 1e-8)
stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0
return stretch
def remove_background_faces(faces, foreground_mask):
"""Keep faces with at least one foreground vertex."""
foreground_points = foreground_mask.reshape(-1)
return faces[foreground_points[faces].any(axis=1)]
def depth_edge_mask(depth, threshold=12.0 / 255.0):
import numpy as np
values = np.asarray(depth, dtype=np.float32)
edges = np.zeros(values.shape, dtype=bool)
for axis in (0, 1):
first = np.take(values, range(values.shape[axis] - 1), axis=axis)
second = np.take(values, range(1, values.shape[axis]), axis=axis)
valid = (
np.isfinite(first)
& np.isfinite(second)
& (first > 0.0)
& (second > 0.0)
)
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
discontinuity = valid & (np.abs(first - second) / scale > threshold)
lower = [slice(None), slice(None)]
upper = [slice(None), slice(None)]
lower[axis] = slice(0, -1)
upper[axis] = slice(1, None)
edges[tuple(lower)] |= discontinuity
edges[tuple(upper)] |= discontinuity
return edges
def edge_culling_stretch_limit(culling):
factor = min(max(float(culling) / 100.0, 0.0001), 1.0)
return 1.0 / factor
def foreground_mesh_mask(foreground_path, depth_shape):
import numpy as np
from PIL import Image, ImageFilter
height, width = depth_shape
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize(
(width, height),
resample=Image.Resampling.BILINEAR,
)
alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75))
return np.asarray(alpha) > 0
def confidence_percentile_threshold(depth, confidence, percentile):
import numpy as np
if float(percentile) <= 0.0:
return None
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,
edge_threshold=12.0 / 255.0,
edge_stretch_limit=None,
vertex_mask=None,
):
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)
)
if confidence_threshold is not None:
valid &= np.isfinite(confidence) & (
confidence >= confidence_threshold
)
if filter_edges:
valid &= ~depth_edge_mask(depth, edge_threshold)
if vertex_mask is not None:
if np.shape(vertex_mask) != depth.shape:
raise ValueError("Foreground mask and depth dimensions do not match")
valid &= np.asarray(vertex_mask, dtype=bool)
triangles = np.concatenate(
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
axis=0,
)
face_valid = valid.reshape(-1)[triangles].all(axis=1)
if edge_stretch_limit is not None:
face_valid &= triangle_edge_stretch_mask(
points,
triangles,
edge_stretch_limit,
)
if max_edge_length > 0.0:
flat_points = points.reshape(-1, 3)
triangle_points = flat_points[triangles]
next_points = np.roll(triangle_points, -1, axis=1)
edge_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
triangles = triangles[face_valid]
if not len(triangles):
return np.empty((0, 3), dtype=np.int32)
return triangles
def triangle_edge_stretch_mask(points, triangles, stretch_limit):
stretch = triangle_edge_stretch_values(points, triangles)
tolerance = float(stretch_limit) * (1.0 + 1e-5)
return (stretch <= tolerance).all(axis=1)

31
src/snapform/server/geometry/projection.py Normal file
View File

@@ -0,0 +1,31 @@
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 unproject_camera_depth(depth, intrinsics):
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
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)

133
src/snapform/server/geometry/surface.py Normal file
View File

@@ -0,0 +1,133 @@
from dataclasses import dataclass, field
from .grid_mesh import (
confidence_percentile_map,
depth_discontinuity_ratio,
mesh_triangle_faces,
remove_background_faces,
)
@dataclass
class GeometryFrame:
depth: object
points: object
intrinsics: object
valid_mask: object
confidence: object | None = None
model_mask: object | None = None
@dataclass
class SurfaceMesh:
vertices: object
faces: object
uv: object
intrinsics: object
image_size: tuple[int, int]
reference_depth: float
point_attributes: dict = field(default_factory=dict)
face_attributes: dict = field(default_factory=dict)
def _validate_frame(frame):
import numpy as np
depth = np.asarray(frame.depth, dtype=np.float32)
points = np.asarray(frame.points, dtype=np.float32)
intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
valid_mask = np.asarray(frame.valid_mask, dtype=bool)
if depth.ndim != 2:
raise ValueError(f"Unsupported depth shape: {depth.shape}")
if points.shape != (*depth.shape, 3):
raise ValueError("Point grid and depth dimensions do not match")
if intrinsics.shape != (3, 3):
raise ValueError(f"Unsupported intrinsics shape: {intrinsics.shape}")
if valid_mask.shape != depth.shape:
raise ValueError("Valid mask and depth dimensions do not match")
if frame.confidence is not None:
confidence = np.asarray(frame.confidence)
if confidence.shape != depth.shape:
raise ValueError("Confidence and depth dimensions do not match")
if frame.model_mask is not None:
model_mask = np.asarray(frame.model_mask)
if model_mask.shape != depth.shape:
raise ValueError("Model mask and depth dimensions do not match")
return depth, points, intrinsics, valid_mask
def _compact_grid(points, faces, depth):
import numpy as np
height, width = depth.shape
used = np.unique(faces)
index_map = np.full(height * width, -1, dtype=np.int32)
index_map[used] = np.arange(len(used), dtype=np.int32)
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
uv = np.stack(np.meshgrid(u, v), axis=-1).reshape(-1, 2)[used]
return (
points.reshape(-1, 3)[used].astype(np.float32),
index_map[faces],
uv,
used,
)
def build_surface(frame, foreground_mask=None):
import numpy as np
depth, points, intrinsics, valid_mask = _validate_frame(frame)
faces = mesh_triangle_faces(
points,
depth,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=valid_mask,
)
if foreground_mask is not None:
foreground_mask = np.asarray(foreground_mask, dtype=bool)
if foreground_mask.shape != depth.shape:
raise ValueError("Foreground mask and depth dimensions do not match")
faces = remove_background_faces(faces, foreground_mask)
if not len(faces):
raise ValueError("Model output does not contain any valid surface faces")
face_attributes = {
"face_depth_discontinuity": (
depth_discontinuity_ratio(depth).reshape(-1)[faces].max(axis=1)
).astype(np.float32),
}
if frame.confidence is not None:
confidence = confidence_percentile_map(
depth,
np.asarray(frame.confidence, dtype=np.float32),
)
face_attributes["face_confidence"] = (
confidence.reshape(-1)[faces].min(axis=1).astype(np.float32)
)
if frame.model_mask is not None:
model_mask = np.asarray(frame.model_mask, dtype=bool)
face_attributes["face_mask"] = model_mask.reshape(-1)[faces].all(axis=1)
vertices, compact_faces, uv, used = _compact_grid(points, faces, depth)
point_depth = depth.reshape(-1)[used].astype(np.float32)
reference_values = point_depth
if frame.model_mask is not None:
masked = point_depth[
np.asarray(frame.model_mask, dtype=bool).reshape(-1)[used]
]
if len(masked):
reference_values = masked
return SurfaceMesh(
vertices=vertices,
faces=compact_faces.astype(np.int32),
uv=uv.astype(np.float32),
intrinsics=intrinsics.astype(np.float32),
image_size=(depth.shape[1], depth.shape[0]),
reference_depth=float(np.median(reference_values)),
point_attributes={"point_depth": point_depth},
face_attributes=face_attributes,
)

5
src/snapform/server/jobs/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
from .depth import run as generate_depth
from .surface import run as generate_surface
__all__ = ("generate_depth", "generate_surface")

78
src/snapform/server/jobs/background.py Normal file
View File

@@ -0,0 +1,78 @@
import shutil
import time
from ..media import copy_frames, prepare_background_input
from ..models.ben2 import infer_foreground, infer_foreground_frames
def run(context, parameters):
"""Run BEN2 and write transparent foreground images into the Job directory."""
total_started = time.perf_counter()
context.progress(0.02, "Starting Environment")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
device = parameters["device"]
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
video_frames = parameters.get("video_frames", 16)
start_frame = parameters.get("start_frame", 1)
images, temporary_directory = prepare_background_input(
parameters["input"],
video_frames,
output_directory,
start_frame=start_frame,
cancel_check=context.check_cancelled,
)
try:
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
context.progress(0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
else:
load_ms, inference_ms = infer_foreground_frames(
images,
model_directory,
device,
output_directory,
context.progress,
0.12,
0.7,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
context.check_cancelled()
total_ms = (time.perf_counter() - total_started) * 1000.0
frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
message = (
f"Background removed{frame_text} in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
)
print(message)
context.progress(1.0, message)
if tuple(output_directory.glob("foreground.*.png")):
return {"foreground": "."}
return {"foreground": "foreground.png"}

20
src/snapform/server/jobs/common.py Normal file
View File

@@ -0,0 +1,20 @@
import shutil
from contextlib import contextmanager
from ..media import prepare_input
@contextmanager
def input_frames(context, parameters):
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
context.directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
yield images
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

26
src/snapform/server/jobs/debug_ben2.py Normal file
View File

@@ -0,0 +1,26 @@
from ..media import copy_frames
from ..models.ben2 import infer_foreground
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing BEN2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
foreground, _load_ms, _inference_ms = infer_foreground(
image,
model_directory,
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
foreground.save(context.directory / f"foreground{suffix}.png")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw BEN2 API foreground")
return {"directory": "."}

21
src/snapform/server/jobs/debug_da3.py Normal file
View File

@@ -0,0 +1,21 @@
from ..media import copy_frames
from ..models import da3
from ..outputs.debug_da3 import write_prediction
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing DA3 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
prediction = da3.infer(
context.resource("model_manager"),
parameters,
images,
cancel_check=context.check_cancelled,
)
write_prediction(prediction, context.directory / "prediction.npz")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw DA3 API prediction")
return {"directory": "."}

26
src/snapform/server/jobs/debug_moge2.py Normal file
View File

@@ -0,0 +1,26 @@
from ..media import copy_frames
from ..models.moge2 import infer_one
from ..outputs.debug_moge2 import write_prediction
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing MoGe-2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_one(
context.resource("model_manager"),
parameters,
image,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
write_prediction(
prediction,
context.directory / f"prediction{suffix}.npz",
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw MoGe-2 API prediction")
return {"directory": "."}

49
src/snapform/server/jobs/depth.py Normal file
View File

@@ -0,0 +1,49 @@
from ..media import copy_frames
from ..models import da3, geometry_adapter, moge2
from ..outputs.depth import write_da3_depth, write_moge_depth
from .common import input_frames
from .foreground import write_optional
def run(context, parameters):
context.progress(0.02, "Preparing depth input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
with input_frames(context, parameters) as images:
if adapter is da3:
prediction = adapter.infer(
model_manager,
parameters,
images,
progress_callback=lambda index, total: context.progress(
0.15 + 0.65 * index / total,
f"Estimating DA3 depth ({index + 1}/{total})",
),
cancel_check=context.check_cancelled,
)
write_da3_depth(prediction, context.directory / "results.npz")
else:
prediction = adapter.infer(
model_manager,
parameters,
images,
include_points=False,
progress_callback=lambda index, total: context.progress(
0.15 + 0.65 * index / total,
f"Estimating MoGe-2 depth ({index + 1}/{total})",
),
cancel_check=context.check_cancelled,
)
write_moge_depth(context.directory, prediction)
write_optional(
context,
parameters,
images,
progress_start=0.84,
progress_span=0.08,
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, f"Exported depth for {len(images)} frame(s)")
return {"depth": "results.npz"}

51
src/snapform/server/jobs/foreground.py Normal file
View File

@@ -0,0 +1,51 @@
from ..models.ben2 import infer_foreground, infer_foreground_frames
def write_optional(
context,
parameters,
images,
*,
progress_start,
progress_span,
):
model_key = parameters.get("remove_background_model", "")
if not model_key:
return []
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(model_key)
device = parameters["device"]
if len(images) == 1:
context.progress(progress_start, "Removing image background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
path = context.directory / "foreground.png"
foreground.convert("RGBA").save(
path,
format="PNG",
compress_level=1,
)
paths = [path]
else:
load_ms, inference_ms = infer_foreground_frames(
images,
model_directory,
device,
context.directory,
context.progress,
progress_start,
progress_span,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
paths = sorted(context.directory.glob("foreground.*.png"))
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)
return paths

21
src/snapform/server/jobs/frames.py Normal file
View File

@@ -0,0 +1,21 @@
import shutil
from ..media import copy_frames, prepare_input
def run(context, parameters):
context.directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
context.directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
copy_frames(images, context.directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
context.progress(1.0, f"Extracted {len(images)} frame(s)")
return {"frames": len(images)}

75
src/snapform/server/jobs/surface.py Normal file
View File

@@ -0,0 +1,75 @@
from ..geometry import foreground_mesh_mask
from ..models import geometry_adapter, moge2
from ..geometry import build_surface
from ..outputs.surface import (
write_normal_texture,
write_surface_archive,
write_surface_texture,
)
from .common import input_frames
from .foreground import write_optional
def run(context, parameters):
context.progress(0.02, "Preparing surface input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Convert to Surface currently supports one image")
context.progress(0.12, "Estimating surface geometry")
if adapter is moge2:
prediction = adapter.infer(
model_manager,
parameters,
images,
include_points=True,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
f"Estimating MoGe-2 geometry ({index + 1}/{total})",
),
cancel_check=context.check_cancelled,
)
else:
prediction = adapter.infer(
model_manager,
parameters,
images,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
f"Estimating DA3 geometry ({index + 1}/{total})",
),
cancel_check=context.check_cancelled,
)
context.check_cancelled()
frame = adapter.surface_frame(prediction)
foreground_paths = write_optional(
context,
parameters,
images,
progress_start=0.72,
progress_span=0.08,
)
foreground_mask = None
if foreground_paths:
foreground_mask = foreground_mesh_mask(
foreground_paths[0],
frame.depth.shape,
)
context.progress(0.82, "Building surface")
mesh = build_surface(frame, foreground_mask=foreground_mask)
context.progress(0.92, "Writing surface")
write_surface_archive(mesh, context.directory / "mesh.npz")
write_surface_texture(
foreground_paths[0] if foreground_paths else images[0],
context.directory / "texture.png",
)
if adapter is moge2:
write_normal_texture(
prediction.normal[0],
context.directory / "normal.png",
)
context.check_cancelled()
context.progress(1.0, "Surface is ready")
return {"surface": "mesh.npz"}

306
src/snapform/server/media.py
View File

@@ -1,306 +0,0 @@
import shutil
import tempfile
import time
from pathlib import Path
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = {
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
}
SCENE_FILENAME = "scene.glb"
class TaskCancelled(RuntimeError):
pass
def check_cancelled(cancel_check):
if cancel_check is not None and cancel_check():
raise TaskCancelled("Task cancelled")
def collect_images(directory):
directory = Path(directory)
return sorted(
path
for path in directory.iterdir()
if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
)
def first_frames(images, maximum_frames, start_frame=1):
"""Return up to maximum_frames image paths in order from start_frame."""
if maximum_frames <= 0:
return images
start_index = max(start_frame - 1, 0)
return images[start_index:start_index + maximum_frames]
def extract_video_frames(
video_path,
output_directory,
maximum_frames,
av,
start_frame=1,
cancel_check=None,
):
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
try:
container = av.open(str(video_path))
except Exception as error:
raise RuntimeError(f"Unable to open video: {video_path}") from error
image_paths = []
skipped = 0
try:
stream = container.streams.video[0]
for frame in container.decode(stream):
check_cancelled(cancel_check)
if skipped < start_frame - 1:
skipped += 1
continue
image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
frame.to_image().save(image_path, format="JPEG")
image_paths.append(image_path)
if len(image_paths) >= maximum_frames:
break
finally:
container.close()
if not image_paths:
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
return image_paths
def prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
check_cancelled(cancel_check)
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
suffix = input_path.suffix.lower()
if suffix in IMAGE_SUFFIXES:
return [input_path], None
if suffix not in VIDEO_SUFFIXES:
supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
raise RuntimeError(f"Unsupported input type; expected one of: {supported}")
import av
temporary_directory = Path(
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
)
try:
images = extract_video_frames(
input_path,
temporary_directory,
video_frames,
av,
start_frame=start_frame,
cancel_check=cancel_check,
)
except Exception:
shutil.rmtree(temporary_directory, ignore_errors=True)
raise
return images, temporary_directory
def prepare_background_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
"""Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
if input_path.suffix.lower() in IMAGE_SUFFIXES:
return [input_path], None
return prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=start_frame,
cancel_check=cancel_check,
)
def copy_frames(images, output_directory):
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
if len(images) <= 1:
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)
def extract_frames(context, parameters):
"""Copy selected source frames into the current job directory."""
context.directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
context.directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
copy_frames(images, context.directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
message = f"Extracted {len(images)} frame(s)"
context.progress(1.0, message)
return {"frames": len(images)}
def exported_paths(output_directory, export_format):
output_directory = Path(output_directory)
if export_format == "glb":
scene_path = output_directory / SCENE_FILENAME
return [scene_path] if scene_path.is_file() else []
if export_format == "gs_ply":
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 == "mesh_data":
result_path = output_directory / "mesh.npz"
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():
return [result_path]
result_path = (
output_directory / "exports" / "mini_npz" / "results.npz"
)
return [result_path] if result_path.is_file() else []
return []
def npz_is_complete(path):
"""Return True only when the npz archive can be fully read."""
try:
import numpy as np
with np.load(path, allow_pickle=False) as archive:
return "depth" in archive.files
except Exception:
return False
def wait_for_exported_paths(
output_directory,
export_format,
timeout=30.0,
cancel_check=None,
):
# DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
# 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
check_cancelled(cancel_check)
result_paths = exported_paths(output_directory, export_format)
if result_paths:
if (
export_format in ("mini_npz", "npz")
and not npz_is_complete(result_paths[0])
):
time.sleep(0.1)
continue
return result_paths
time.sleep(0.05)
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

26
src/snapform/server/media/__init__.py Normal file
View File

@@ -0,0 +1,26 @@
from .input import (
IMAGE_SUFFIXES,
VIDEO_SUFFIXES,
TaskCancelled,
check_cancelled,
collect_images,
copy_frames,
extract_video_frames,
first_frames,
prepare_background_input,
prepare_input,
)
__all__ = (
"IMAGE_SUFFIXES",
"VIDEO_SUFFIXES",
"TaskCancelled",
"check_cancelled",
"collect_images",
"copy_frames",
"extract_video_frames",
"first_frames",
"prepare_background_input",
"prepare_input",
)

158
src/snapform/server/media/input.py Normal file
View File

@@ -0,0 +1,158 @@
import shutil
import tempfile
from pathlib import Path
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = {
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
}
class TaskCancelled(RuntimeError):
pass
def check_cancelled(cancel_check):
if cancel_check is not None and cancel_check():
raise TaskCancelled("Task cancelled")
def collect_images(directory):
directory = Path(directory)
return sorted(
path
for path in directory.iterdir()
if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
)
def first_frames(images, maximum_frames, start_frame=1):
"""Return up to maximum_frames image paths in order from start_frame."""
if maximum_frames <= 0:
return images
start_index = max(start_frame - 1, 0)
return images[start_index:start_index + maximum_frames]
def extract_video_frames(
video_path,
output_directory,
maximum_frames,
av,
start_frame=1,
cancel_check=None,
):
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
try:
container = av.open(str(video_path))
except Exception as error:
raise RuntimeError(f"Unable to open video: {video_path}") from error
image_paths = []
skipped = 0
try:
stream = container.streams.video[0]
for frame in container.decode(stream):
check_cancelled(cancel_check)
if skipped < start_frame - 1:
skipped += 1
continue
image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
frame.to_image().save(image_path, format="JPEG")
image_paths.append(image_path)
if len(image_paths) >= maximum_frames:
break
finally:
container.close()
if not image_paths:
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
return image_paths
def prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
check_cancelled(cancel_check)
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
suffix = input_path.suffix.lower()
if suffix in IMAGE_SUFFIXES:
return [input_path], None
if suffix not in VIDEO_SUFFIXES:
supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
raise RuntimeError(f"Unsupported input type; expected one of: {supported}")
import av
temporary_directory = Path(
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
)
try:
images = extract_video_frames(
input_path,
temporary_directory,
video_frames,
av,
start_frame=start_frame,
cancel_check=cancel_check,
)
except Exception:
shutil.rmtree(temporary_directory, ignore_errors=True)
raise
return images, temporary_directory
def prepare_background_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
"""Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
if input_path.suffix.lower() in IMAGE_SUFFIXES:
return [input_path], None
return prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=start_frame,
cancel_check=cancel_check,
)
def copy_frames(images, output_directory):
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
if len(images) <= 1:
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)

6
src/snapform/server/model_manager.py
View File

@@ -27,7 +27,7 @@ class ModelSessionCache:
return self._load_ben2(model_dir, requested_device, key)
def _load_ben2(self, model_dir, requested_device, key):
from .onnx_ben2 import create_session
from .models.onnx_ben2 import create_session
started = time.perf_counter()
model = create_session(model_dir, requested_device)
@@ -47,7 +47,7 @@ class ModelSessionCache:
return self._load_da3(model_dir, requested_device, key)
def _load_da3(self, model_dir, requested_device, key):
from .onnx_da3 import create_session
from .models.onnx_da3 import create_session
model = create_session(model_dir, requested_device)
self.da3_model = model
@@ -55,7 +55,7 @@ class ModelSessionCache:
return model
def get_moge2(self, model_dir, requested_device):
from .onnx_moge2 import create_session
from .models.onnx_moge2 import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None:

8
src/snapform/server/models/__init__.py Normal file
View File

@@ -0,0 +1,8 @@
from . import ben2, da3, moge2
def geometry_adapter(model_key):
return moge2 if model_key.startswith("MOGE2_") else da3
__all__ = ("ben2", "da3", "moge2", "geometry_adapter")

78
src/snapform/server/models/ben2.py Normal file
View File

@@ -0,0 +1,78 @@
import time
from pathlib import Path
from ..media import IMAGE_SUFFIXES, check_cancelled
def infer_foreground(
input_path,
model_dir,
device,
model_cache=None,
cancel_check=None,
):
from PIL import Image
input_path = Path(input_path)
check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
from .onnx_ben2 import create_session, infer
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
if owns_model:
load_started = time.perf_counter()
model = create_session(model_dir, requested_device)
load_ms = (time.perf_counter() - load_started) * 1000.0
else:
model, load_ms = model_cache.get_ben2(model_dir, device)
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
inference_started = time.perf_counter()
foreground = infer(model, image)
inference_ms = (time.perf_counter() - inference_started) * 1000.0
check_cancelled(cancel_check)
return foreground, load_ms, inference_ms
def infer_foreground_frames(
images,
model_dir,
device,
output_directory,
progress_callback,
start_progress,
progress_span,
model_cache=None,
cancel_check=None,
):
"""Run BEN2 on every frame and write foreground.0001.png etc."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
progress_callback(
start_progress + progress_span * index / len(images),
f"Removing background ({index + 1}/{len(images)})",
)
foreground, load_ms, inference_ms = infer_foreground(
image_path,
model_dir,
device,
model_cache=model_cache,
cancel_check=cancel_check,
)
output_path = output_directory / f"foreground.{index + 1:04d}.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError(f"BEN2 did not write: {output_path}")
total_load_ms += load_ms
total_inference_ms += inference_ms
return total_load_ms, total_inference_ms

94
src/snapform/server/models/da3.py Normal file
View File

@@ -0,0 +1,94 @@
from ..geometry import unproject_camera_depth
from ..geometry import GeometryFrame
def merge_frame_predictions(predictions):
import numpy as np
if not predictions:
raise ValueError("DA3 returned no frame predictions")
merged = predictions[0]
for name in (
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
):
setattr(
merged,
name,
np.concatenate(
[getattr(prediction, name) for prediction in predictions],
axis=0,
),
)
return merged
def infer(
model_manager,
parameters,
images,
*,
progress_callback=None,
cancel_check=None,
):
from .onnx_da3 import infer as infer_onnx
from .onnx_runtime import provider_label
model_directory = model_manager.directory(parameters["model"])
device = parameters["device"]
model = model_manager.get_da3(model_directory, device)
options = {
"process_res": parameters["process_res"],
"process_res_method": parameters["process_res_method"],
"ref_view_strategy": parameters["ref_view_strategy"],
}
def run(paths):
try:
return infer_onnx(model, paths, **options)
except UnicodeDecodeError as error:
if provider_label(model) == "DirectML":
raise RuntimeError(
"DA3 inference failed on DirectML. Switch the device to "
"CPU in SnapForm Preferences and retry. "
f"Original error: {error}"
) from error
raise
if parameters.get("inference_mode", "joint") != "per_frame":
return run([str(path) for path in images])
predictions = []
for index, image in enumerate(images):
if cancel_check is not None:
cancel_check()
if progress_callback is not None:
progress_callback(index, len(images))
predictions.append(run([str(image)]))
return merge_frame_predictions(predictions)
def surface_frame(prediction):
import numpy as np
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
intrinsics = np.asarray(prediction.intrinsics, dtype=np.float32)
if depth.shape[0] != 1:
raise ValueError("Convert to Surface currently supports one image")
points = unproject_camera_depth(depth[0], intrinsics[0])
valid = (
np.isfinite(depth[0])
& (depth[0] > 0.0)
& np.isfinite(points).all(axis=-1)
)
return GeometryFrame(
depth=depth[0],
points=points,
intrinsics=intrinsics[0],
valid_mask=valid,
confidence=confidence[0],
)

119
src/snapform/server/models/moge2.py Normal file
View File

@@ -0,0 +1,119 @@
from dataclasses import dataclass
from ..geometry import GeometryFrame
@dataclass
class Prediction:
depth: object
normal: object
mask: object
intrinsics: object
points: object | None
def archive_values(self):
values = {
"depth": self.depth,
"normal": self.normal,
"mask": self.mask,
"intrinsics": self.intrinsics,
}
if self.points is not None:
values["points"] = self.points
return values
def infer(
model_manager,
parameters,
images,
*,
include_points,
progress_callback=None,
cancel_check=None,
):
import numpy as np
depths = []
normals = []
masks = []
intrinsics = []
points = []
for index, image_path in enumerate(images):
if cancel_check is not None:
cancel_check()
if progress_callback is not None:
progress_callback(index, len(images))
prediction = infer_one(
model_manager,
parameters,
image_path,
)
depths.append(prediction["depth"])
normals.append(prediction["normal"])
masks.append(prediction["mask"].astype(bool))
intrinsics.append(prediction["intrinsics"])
if include_points:
points.append(prediction["points"])
if len({array.shape for array in depths}) != 1:
raise RuntimeError("MoGe-2 frame sizes must match")
return Prediction(
depth=np.stack(depths).astype(np.float32),
normal=np.stack(normals).astype(np.float32),
mask=np.stack(masks),
intrinsics=np.stack(intrinsics).astype(np.float32),
points=(
np.stack(points).astype(np.float32)
if include_points
else None
),
)
def infer_one(model_manager, parameters, image_path):
import numpy as np
from PIL import Image
from .onnx_moge2 import infer as infer_onnx
model_directory = model_manager.directory(parameters["model"])
model = model_manager.get_moge2(
model_directory,
parameters["device"],
)
with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
return infer_onnx(
model,
image,
int(parameters["resolution_level"]),
)
def surface_frame(prediction):
import numpy as np
if prediction.depth.shape[0] != 1 or prediction.points is None:
raise ValueError("Convert to Surface requires one MoGe-2 point grid")
depth = prediction.depth[0]
points = prediction.points[0]
intrinsics = prediction.intrinsics[0].copy()
height, width = depth.shape
if (
0.0 <= float(intrinsics[0, 2]) <= 1.0
and 0.0 <= float(intrinsics[1, 2]) <= 1.0
):
intrinsics[0] *= width
intrinsics[1] *= height
intrinsics[2] = (0.0, 0.0, 1.0)
valid = (
np.isfinite(depth)
& (depth > 0.0)
& np.isfinite(points).all(axis=-1)
)
return GeometryFrame(
depth=depth,
points=points,
intrinsics=intrinsics,
valid_mask=valid,
model_mask=prediction.mask[0],
)

0
src/snapform/server/onnx_ben2.py → src/snapform/server/models/onnx_ben2.py
View File

0
src/snapform/server/onnx_da3.py → src/snapform/server/models/onnx_da3.py
View File

0
src/snapform/server/onnx_moge2.py → src/snapform/server/models/onnx_moge2.py
View File

0
src/snapform/server/onnx_runtime.py → src/snapform/server/models/onnx_runtime.py
View File

304
src/snapform/server/moge2.py
View File

@@ -1,304 +0,0 @@
import shutil
from pathlib import Path
from .background import infer_foreground, infer_foreground_frames
from .geometry import (
depth_discontinuity_ratio,
depth_edge_mask,
foreground_mesh_mask,
mesh_triangle_faces,
)
from .media import copy_frames, prepare_input
def _save_moge_preview(output_directory, index, depth, normal, mask):
import numpy as np
from PIL import Image
suffix = "" if index is None else f".{index + 1:04d}"
valid = mask & np.isfinite(depth) & (depth > 0)
preview = np.zeros(depth.shape, dtype=np.uint8)
if np.any(valid):
low, high = np.percentile(depth[valid], (2.0, 98.0))
span = max(float(high - low), 1e-6)
preview[valid] = np.clip(
(depth[valid] - low) / span * 255.0,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(preview, mode="L").save(
output_directory / f"depth{suffix}.png"
)
normal_rgb = np.clip(
(normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0)
* 127.5,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(normal_rgb, mode="RGB").save(
output_directory / f"normal{suffix}.png"
)
Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
output_directory / f"mask{suffix}.png"
)
def export_moge_debug_mesh(
output_directory,
points,
depth,
mask,
normal,
image,
export_format,
):
import numpy as np
import trimesh
clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04)
faces = mesh_triangle_faces(
points,
depth,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=clean_mask,
)
vertices = points.reshape(-1, 3)
colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8)
vertex_normals = normal.reshape(-1, 3)
vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32)
vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
mesh = trimesh.Trimesh(
vertices=vertices,
faces=faces,
vertex_colors=colors,
vertex_normals=vertex_normals,
process=False,
)
mesh.remove_unreferenced_vertices()
filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply"
mesh.export(Path(output_directory) / filename)
def export_moge_surface_archive(
output_path,
points,
depth,
mask,
normal,
intrinsics,
foreground_mask=None,
):
import numpy as np
height, width = depth.shape
finite_mask = np.isfinite(depth) & (depth > 0.0)
finite_mask &= np.isfinite(points).all(axis=-1)
if foreground_mask is not None:
finite_mask &= foreground_mask
depth_discontinuity = depth_discontinuity_ratio(depth)
faces = mesh_triangle_faces(
points,
depth,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=finite_mask,
)
vertices = points.reshape(-1, 3)
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)
vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2)
pixel_intrinsics = intrinsics.copy()
pixel_intrinsics[0] *= width
pixel_intrinsics[1] *= height
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = depth.reshape(-1).astype(np.float32)
face_mask = mask.reshape(-1)[faces].all(axis=1)
face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
masked_depth = point_depth[mask.reshape(-1).astype(bool)]
reference_depth = np.median(masked_depth if len(masked_depth) else point_depth)
np.savez(
output_path,
vertices=vertices.astype(np.float32),
faces=faces.astype(np.int32),
uv=vertex_uvs.astype(np.float32),
intrinsics=pixel_intrinsics.astype(np.float32),
image_size=np.array((width, height), dtype=np.int32),
reference_depth=np.float32(reference_depth),
point_depth=point_depth,
normal_map=normal.astype(np.float32),
face_mask=face_mask,
face_depth_discontinuity=face_depth_discontinuity,
)
def generate_moge2(context, parameters):
"""Run MoGe-2 inference and return the artifact used by the requesting Job."""
import numpy as np
from PIL import Image
from .onnx_moge2 import (
infer as infer_moge2,
provider_label,
)
context.progress(0.02, "Starting MoGe-2")
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
device = parameters["device"]
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
output_directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
model = model_manager.get_moge2(model_directory, device)
context.progress(
0.08,
f"Loading MoGe-2 on {provider_label(model)}",
)
depths = []
normals = []
masks = []
points = []
intrinsics = []
source_images = []
for index, image_path in enumerate(images):
context.check_cancelled()
context.progress(
0.15 + 0.68 * index / len(images),
f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})",
)
with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
source_images.append(image)
prediction = infer_moge2(
model,
image,
int(parameters["resolution_level"]),
)
depth = prediction["depth"]
normal = prediction["normal"]
mask = prediction["mask"].astype(bool)
point = prediction["points"]
intrinsic = prediction["intrinsics"]
depths.append(depth)
normals.append(normal)
masks.append(mask)
intrinsics.append(intrinsic)
if bool(
parameters.get("include_points", False)
or parameters.get("surface_mesh", False)
):
points.append(point)
debug_format = parameters.get("export_format", "")
if not debug_format or debug_format == "maps":
_save_moge_preview(
output_directory,
None if len(images) == 1 else index,
depth,
normal,
mask,
)
shapes = {array.shape for array in depths}
if len(shapes) != 1:
raise RuntimeError("MoGe-2 frame sizes must match")
result = {
"depth": np.stack(depths).astype(np.float32),
"normal": np.stack(normals).astype(np.float32),
"mask": np.stack(masks),
"intrinsics": np.stack(intrinsics).astype(np.float32),
}
if points:
result["points"] = np.stack(points).astype(np.float32)
debug_format = parameters.get("export_format", "")
if debug_format in {"glb", "ply"} and len(images) == 1:
context.progress(
0.84,
f"Exporting MoGe-2 {debug_format.upper()}",
)
export_moge_debug_mesh(
output_directory,
points[0],
depths[0],
masks[0],
normals[0],
source_images[0],
debug_format,
)
foreground_mask = None
remove_background_model = parameters.get("remove_background_model", "")
if remove_background_model:
remove_background_model_dir = model_manager.directory(
remove_background_model
)
context.progress( 0.87, "Removing image background")
if len(images) == 1:
foreground, _, _ = infer_foreground(
images[0],
remove_background_model_dir,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
foreground_path = output_directory / "foreground.png"
foreground.save(foreground_path)
if points:
foreground_mask = foreground_mesh_mask(
foreground_path,
result["depth"][0].shape,
)
else:
infer_foreground_frames(
images,
remove_background_model_dir,
device,
output_directory,
context.progress,
0.87,
0.08,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
if bool(parameters.get("surface_mesh", False)):
if len(images) != 1 or not points:
raise RuntimeError("MoGe-2 Surface requires one point map")
context.progress( 0.94, "Building official MoGe-2 surface")
export_moge_surface_archive(
output_directory / "mesh.npz",
points[0],
result["depth"][0],
result["mask"][0],
result["normal"][0],
result["intrinsics"][0],
foreground_mask=(
foreground_mask
if remove_background_model
else None
),
)
np.savez(output_directory / "results.npz", **result)
copy_frames(images, output_directory)
context.progress(
1.0,
f"Exported MoGe-2 data for {len(images)} frame(s)",
)
if context.job_type == "generate-moge2-surface":
return {"surface": "mesh.npz"}
if context.job_type == "generate-moge2-depth":
return {"depth": "results.npz"}
return {"directory": "."}
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

11
src/snapform/server/outputs/__init__.py Normal file
View File

@@ -0,0 +1,11 @@
from .depth import write_da3_depth, write_moge_depth
from .surface import write_normal_texture, write_surface_archive, write_surface_texture
__all__ = (
"write_da3_depth",
"write_moge_depth",
"write_normal_texture",
"write_surface_archive",
"write_surface_texture",
)

24
src/snapform/server/outputs/debug_da3.py Normal file
View File

@@ -0,0 +1,24 @@
from pathlib import Path
def write_prediction(prediction, output_path):
import numpy as np
values = {}
for name in (
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
"is_metric",
"sky",
"scale_factor",
):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.asarray(value)
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez(output_path, **values)
return output_path

13
src/snapform/server/outputs/debug_moge2.py Normal file
View File

@@ -0,0 +1,13 @@
from pathlib import Path
def write_prediction(prediction, output_path):
import numpy as np
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez(
output_path,
**{name: np.asarray(value) for name, value in prediction.items()},
)
return output_path

64
src/snapform/server/outputs/depth.py Normal file
View File

@@ -0,0 +1,64 @@
from pathlib import Path
from .surface import write_normal_texture
def write_da3_depth(prediction, output_path):
import numpy as np
values = {"depth": np.round(prediction.depth, 8)}
for name in ("conf", "extrinsics", "intrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.round(value, 2) if name == "conf" else value
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output_path, **values)
return output_path
def write_moge_depth(output_directory, prediction):
import numpy as np
output_directory = Path(output_directory)
np.savez(
output_directory / "results.npz",
**prediction.archive_values(),
)
for index, (depth, normal, mask) in enumerate(
zip(prediction.depth, prediction.normal, prediction.mask)
):
_write_moge_maps(
output_directory,
None if len(prediction.depth) == 1 else index,
depth,
normal,
mask,
)
def _write_moge_maps(output_directory, index, depth, normal, mask):
import numpy as np
from PIL import Image
suffix = "" if index is None else f".{index + 1:04d}"
valid = mask & np.isfinite(depth) & (depth > 0)
preview = np.zeros(depth.shape, dtype=np.uint8)
if np.any(valid):
low, high = np.percentile(depth[valid], (2.0, 98.0))
span = max(float(high - low), 1e-6)
preview[valid] = np.clip(
(depth[valid] - low) / span * 255.0,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(preview, mode="L").save(
output_directory / f"depth{suffix}.png"
)
write_normal_texture(
normal,
output_directory / f"normal{suffix}.png",
)
Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
output_directory / f"mask{suffix}.png"
)

44
src/snapform/server/outputs/surface.py Normal file
View File

@@ -0,0 +1,44 @@
from pathlib import Path
def write_surface_archive(mesh, output_path):
import numpy as np
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
output_path,
vertices=mesh.vertices,
faces=mesh.faces,
uv=mesh.uv,
intrinsics=mesh.intrinsics,
image_size=np.asarray(mesh.image_size, dtype=np.int32),
reference_depth=np.asarray(mesh.reference_depth, dtype=np.float32),
**mesh.point_attributes,
**mesh.face_attributes,
)
return output_path
def write_surface_texture(image_path, output_path):
from PIL import Image
output_path = Path(output_path)
with Image.open(image_path) as image:
image.convert("RGBA").save(output_path, format="PNG", compress_level=1)
return output_path
def write_normal_texture(normal, output_path):
import numpy as np
from PIL import Image
normal_rgb = np.clip(
(normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0)
* 127.5,
0.0,
255.0,
).astype(np.uint8)
output_path = Path(output_path)
Image.fromarray(normal_rgb, mode="RGB").save(output_path)
return output_path