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

feat: Add MoGe-2 support and enhance depth plane functionality

Browse Source 
- Introduced new MoGe-2 depth plane group with geometry and material displacement options.
- Updated existing depth plane group to improve background depth handling and subdivision settings.
- Enhanced tests to validate MoGe-2 models and ensure proper integration with existing systems.
- Modified installation environment to include new dependencies for MoGe-2 and utils3d.
- Updated vendor lock file to include new artifacts for MoGe and utils3d.
- Refactored download and model handling to accommodate new MoGe-2 models.
- Improved validation scripts for depth plane assets to ensure compliance with new specifications.
This commit is contained in:
马南
2026-08-13 15:55:58 +08:00
parent 92d2cea1d1
commit a6370d678d
25 changed files with 1770 additions and 258 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
src/snapform/operators.py
View File

@@ -1,3 +1,4 @@
import os
import shutil
import subprocess
import threading
@@ -9,6 +10,7 @@ import bpy
from . import client
from .models import (
BEN2_MODEL,
MOGE2_MODEL,
get_downloadable_model,
get_model,
selected_model_key,
@@ -17,6 +19,8 @@ from .models import (
RECONSTRUCTION_JOB = "RECONSTRUCT"
IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
MOGE2_DEPTH_JOB = "MOGE2_DEPTH"
MOGE2_DEBUG_JOB = "MOGE2_DEBUG"
IMAGE_MESH_JOB = "IMAGE_MESH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
@@ -26,6 +30,7 @@ DA3_IMAGE_JOBS = frozenset({IMAGE_DEPTH_JOB, IMAGE_MESH_JOB})
IMAGE_EMPTY_JOBS = frozenset(
{
*DA3_IMAGE_JOBS,
MOGE2_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
@@ -34,6 +39,7 @@ IMAGE_EMPTY_JOBS = frozenset(
PRODUCTION_AI_JOBS = frozenset(
{
*DA3_IMAGE_JOBS,
MOGE2_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}
@@ -156,6 +162,8 @@ class RunJob(bpy.types.Operator):
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(MOGE2_DEPTH_JOB, "Convert Image with MoGe-2", ""),
(MOGE2_DEBUG_JOB, "Test MoGe-2 API", ""),
(IMAGE_MESH_JOB, "Convert Image to Surface", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
@@ -210,6 +218,12 @@ class RunJob(bpy.types.Operator):
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_moge_resolution_level: bpy.props.IntProperty(
default=5,
min=0,
max=9,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
global _active_job
@@ -236,6 +250,8 @@ 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}:
export_format = "moge_npz"
job = ActiveJob(
kind=self.job_type,
@@ -300,12 +316,20 @@ class RunJob(bpy.types.Operator):
if job.kind == "INSTALL_ENV":
client.stop_service()
job.log_file = job.log_path.open("w", encoding="utf-8")
process_environment = None
if job.kind == "INSTALL_ENV":
process_environment = os.environ.copy()
process_environment["PYTHONIOENCODING"] = (
"utf-8:backslashreplace"
)
process_environment["PYTHONUTF8"] = "1"
job.process = subprocess.Popen(
command,
stdout=job.log_file,
stderr=subprocess.STDOUT,
cwd=working_directory,
creationflags=creation_flags,
env=process_environment,
)
job.started = True
if job.cancelled:
@@ -401,6 +425,35 @@ 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
model = get_downloadable_model(model_key)
if not client.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
if remove_background and not client.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
client.moge_request(
settings,
paths["status"],
paths["result"],
input_path=input_path,
resolution_level=(
self.snapform_moge_resolution_level
if self.job_type == MOGE2_DEPTH_JOB
else settings.moge_resolution_level
),
include_points=self.job_type == MOGE2_DEBUG_JOB,
remove_background=(
self.job_type == MOGE2_DEPTH_JOB
and remove_background
),
start_frame=start_frame,
device=self._task_device(settings),
model_key=model_key,
),
paths,
)
if self.job_type not in {
RECONSTRUCTION_JOB,
IMAGE_DEPTH_JOB,
@@ -478,6 +531,8 @@ 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}:
return "Starting MoGe-2 inference"
if self.job_type == IMAGE_MESH_JOB:
return "Starting perspective surface generation"
if self.job_type == BACKGROUND_REMOVAL_JOB:
@@ -626,6 +681,7 @@ class RunJob(bpy.types.Operator):
if job.kind in {
RECONSTRUCTION_JOB,
BACKGROUND_REMOVAL_JOB,
MOGE2_DEBUG_JOB,
}:
try:
open_job_directory(job.result_directory)
@@ -635,8 +691,10 @@ class RunJob(bpy.types.Operator):
message = (
f"Exported {job.export_format} to {job.result_directory}"
)
else:
elif job.kind == BACKGROUND_REMOVAL_JOB:
message = f"Removed background to {job.result_directory}"
else:
message = f"Exported MoGe-2 data to {job.result_directory}"
self.report({"INFO"}, message)
if job.kind == IMAGE_DEPTH_JOB:
try:
@@ -653,6 +711,23 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Depth plane is ready"
if job.kind == MOGE2_DEPTH_JOB:
try:
settings.status = "Building MoGe-2 depth plane"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_moge2_depth_plane(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "MoGe-2 depth 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 depth plane is ready"
if job.kind == IMAGE_MESH_JOB:
try:
settings.status = "Building perspective surface"
@@ -790,6 +865,35 @@ class RunJob(bpy.types.Operator):
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the depth plane Undo step")
def _build_moge2_depth_plane(self, job):
result_directory = job.result_directory
is_sequence = bool(list(result_directory.glob("frame.*")))
foreground_sequence = bool(
list(result_directory.glob("foreground.*.png"))
)
result = bpy.ops.snapform.build_moge2_depth_plane(
snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(result_directory / "results.npz"),
snapform_background_path=(
str(result_directory)
if job.remove_background and foreground_sequence
else str(result_directory / "foreground.png")
if job.remove_background
else ""
),
snapform_color_frames_dir=(
str(result_directory) if is_sequence else ""
),
snapform_data_frames_dir=str(result_directory),
snapform_delete_source=job.delete_source,
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the MoGe-2 depth plane")
undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the depth plane Undo step")
def _apply_removed_background(self, job, push_undo=True):
result_directory = job.result_directory
foreground_paths = sorted(result_directory.glob("foreground.*.png"))