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

63
src/snapform/operators.py
View File

@@ -20,6 +20,7 @@ from .models import (
RECONSTRUCTION_JOB = "RECONSTRUCT"
IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
MOGE2_DEPTH_JOB = "MOGE2_DEPTH"
MOGE2_MESH_JOB = "MOGE2_MESH"
MOGE2_DEBUG_JOB = "MOGE2_DEBUG"
IMAGE_MESH_JOB = "IMAGE_MESH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
@@ -31,6 +32,7 @@ IMAGE_EMPTY_JOBS = frozenset(
{
*DA3_IMAGE_JOBS,
MOGE2_DEPTH_JOB,
MOGE2_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
@@ -40,6 +42,7 @@ PRODUCTION_AI_JOBS = frozenset(
{
*DA3_IMAGE_JOBS,
MOGE2_DEPTH_JOB,
MOGE2_MESH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}
@@ -163,6 +166,7 @@ class RunJob(bpy.types.Operator):
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(MOGE2_DEPTH_JOB, "Convert Image with MoGe-2", ""),
(MOGE2_MESH_JOB, "Convert Image to Surface with MoGe-2", ""),
(MOGE2_DEBUG_JOB, "Test MoGe-2 API", ""),
(IMAGE_MESH_JOB, "Convert Image to Surface", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
@@ -250,7 +254,7 @@ class RunJob(bpy.types.Operator):
export_format = "mini_npz"
elif self.job_type == IMAGE_MESH_JOB:
export_format = "mesh_data"
elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}:
elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}:
export_format = "moge_npz"
job = ActiveJob(
@@ -425,8 +429,12 @@ class RunJob(bpy.types.Operator):
),
paths,
)
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}:
model_key = self.snapform_model_key or MOGE2_MODEL.key
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}:
model_key = (
settings.moge_model
if self.job_type == MOGE2_DEBUG_JOB
else self.snapform_model_key or MOGE2_MODEL.key
)
model = get_downloadable_model(model_key)
if not client.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
@@ -440,17 +448,23 @@ class RunJob(bpy.types.Operator):
input_path=input_path,
resolution_level=(
self.snapform_moge_resolution_level
if self.job_type == MOGE2_DEPTH_JOB
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB}
else settings.moge_resolution_level
),
include_points=self.job_type == MOGE2_DEBUG_JOB,
include_points=self.job_type in {MOGE2_MESH_JOB, MOGE2_DEBUG_JOB},
remove_background=(
self.job_type == MOGE2_DEPTH_JOB
self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB}
and remove_background
),
start_frame=start_frame,
device=self._task_device(settings),
model_key=model_key,
debug_export_format=(
settings.moge_export_format
if self.job_type == MOGE2_DEBUG_JOB
else ""
),
surface_mesh=self.job_type == MOGE2_MESH_JOB,
),
paths,
)
@@ -531,7 +545,7 @@ class RunJob(bpy.types.Operator):
return f"Starting {settings.export_format} export"
if self.job_type == IMAGE_DEPTH_JOB:
return "Starting depth estimation"
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_DEBUG_JOB}:
if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}:
return "Starting MoGe-2 inference"
if self.job_type == IMAGE_MESH_JOB:
return "Starting perspective surface generation"
@@ -728,6 +742,21 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "MoGe-2 depth plane is ready"
if job.kind == MOGE2_MESH_JOB:
try:
settings.status = "Building MoGe-2 surface"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_moge2_image_mesh(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "MoGe-2 surface import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "MoGe-2 surface is ready"
if job.kind == IMAGE_MESH_JOB:
try:
settings.status = "Building perspective surface"
@@ -801,6 +830,26 @@ class RunJob(bpy.types.Operator):
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the surface Undo step")
def _build_moge2_image_mesh(self, job):
result_directory = job.result_directory
mesh_path = result_directory / "mesh.npz"
if not mesh_path.is_file():
raise RuntimeError("MoGe-2 did not export the official surface mesh")
result = bpy.ops.snapform.build_image_mesh(
snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(mesh_path),
snapform_delete_source=job.delete_source,
snapform_create_camera=job.mesh_create_camera,
snapform_unshaded=job.unshaded,
snapform_remove_background=False,
snapform_moge2=True,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the MoGe-2 perspective surface")
undo_result = bpy.ops.ed.undo_push(message="Convert to Surface")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the surface Undo step")
def _build_depth_plane(self, context, job):
result_directory = job.result_directory
npz_path = result_directory / "results.npz"