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

feat: Add MoGe-2 surface generation and export functionality

Browse Source 
- Introduced new enum items for depth plane and surface models with caching.
- Added MOGE2_MESH_JOB to operators for converting images to surfaces.
- Implemented surface mesh building and exporting in RunJob operator.
- Enhanced UI to include export format options for MoGe-2.
- Created validation tools for MoGe-2 surface assets and depth planes.
- Updated geometry nodes to support new surface generation features.
- Refactored existing tests to accommodate new functionality and ensure stability.
This commit is contained in:
马南
2026-08-13 22:07:36 +08:00
parent a6370d678d
commit caeb6d706c
17 changed files with 902 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
src/snapform/worker.py
View File

@@ -1514,6 +1514,94 @@ def _save_moge_preview(output_directory, index, depth, normal, mask):
)
def export_moge_debug_mesh(
output_directory,
points,
depth,
mask,
normal,
image,
export_format,
):
import numpy as np
import utils3d
from moge.utils.io import save_glb, save_ply
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
uv = utils3d.np.uv_map(*depth.shape)
result = utils3d.np.build_mesh_from_map(
points,
image,
uv,
normal,
mask=clean_mask,
tri=True,
)
faces, vertices, colors, vertex_uvs, vertex_normals = result
vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32)
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
if export_format == "glb":
save_glb(
output_directory / "mesh.glb",
vertices,
faces,
vertex_uvs,
(image * 255.0).astype(np.uint8),
vertex_normals,
)
return
if export_format == "ply":
save_ply(
output_directory / "pointcloud.ply",
vertices,
np.zeros((0, 3), dtype=np.int32),
colors,
vertex_normals,
)
def export_moge_surface_archive(
output_path,
points,
depth,
mask,
normal,
intrinsics,
):
import numpy as np
import utils3d
height, width = depth.shape
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
result = utils3d.np.build_mesh_from_map(
points,
utils3d.np.uv_map(height, width),
normal,
mask=clean_mask,
tri=True,
)
faces, vertices, vertex_uvs, vertex_normals = result
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
pixel_intrinsics = intrinsics.copy()
pixel_intrinsics[0] *= width
pixel_intrinsics[1] *= height
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = vertices[:, 2].astype(np.float32)
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(np.median(point_depth)),
point_depth=point_depth,
vertex_normal=vertex_normals.astype(np.float32),
face_mask=np.ones(len(faces), dtype=bool),
)
def run_moge(args, model_cache=None, cancel_check=None):
import numpy as np
from PIL import Image
@@ -1550,6 +1638,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
masks = []
points = []
intrinsics = []
source_images = []
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
write_status(
@@ -1560,6 +1649,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device)
source_images.append(image)
with torch.inference_mode():
prediction = model.infer(
image_tensor,
@@ -1579,13 +1669,15 @@ def run_moge(args, model_cache=None, cancel_check=None):
intrinsics.append(intrinsic)
if bool(getattr(args, "include_points", False)):
points.append(point)
_save_moge_preview(
output_directory,
None if len(images) == 1 else index,
depth,
normal,
mask,
)
debug_format = getattr(args, "debug_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:
@@ -1598,7 +1690,18 @@ def run_moge(args, model_cache=None, cancel_check=None):
}
if points:
result["points"] = np.stack(points).astype(np.float32)
np.savez(output_directory / "results.npz", **result)
debug_format = getattr(args, "debug_export_format", "")
if debug_format in {"glb", "ply"} and len(images) == 1:
write_status(args.status, 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,
)
remove_background_model_dir = getattr(
args, "remove_background_model_dir", ""
@@ -1614,7 +1717,14 @@ def run_moge(args, model_cache=None, cancel_check=None):
model_cache=model_cache,
cancel_check=cancel_check,
)
foreground.save(output_directory / "foreground.png")
foreground_path = output_directory / "foreground.png"
foreground.save(foreground_path)
if points:
foreground_mask = foreground_mesh_mask(
foreground_path,
result["depth"][0].shape,
)
result["mask"][0] &= foreground_mask
else:
infer_foreground_frames(
images,
@@ -1628,6 +1738,19 @@ def run_moge(args, model_cache=None, cancel_check=None):
model_cache=model_cache,
cancel_check=cancel_check,
)
if bool(getattr(args, "surface_mesh", False)):
if len(images) != 1 or not points:
raise RuntimeError("MoGe-2 Surface requires one point map")
write_status(args.status, 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],
)
np.savez(output_directory / "results.npz", **result)
copy_frames(images, output_directory)
write_status(
args.status,