feat: Add MoGe-2 support and enhance depth plane functionality
- 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:
@@ -13,12 +13,14 @@ ORIGINAL_DISPLACEMENT_STRENGTH = 1.0
|
||||
ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0
|
||||
BUMP_BASE_DISTANCE = 0.001
|
||||
DEPTH_PLANE_NODE_GROUP_NAME = "O DA3 Depth Plane"
|
||||
MOGE_DEPTH_PLANE_NODE_GROUP_NAME = "O MoGe2 Depth Plane"
|
||||
NODE_ASSET_NAME = "O_SnapForm.blend"
|
||||
MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement"
|
||||
BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth"
|
||||
DEPTH_SCALE_ATTRIBUTE = "o_depth_scale"
|
||||
SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer"
|
||||
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer"
|
||||
MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer"
|
||||
DA3_PATCH_SIZE = 14
|
||||
DA3_MIN_PROCESS_RES = 252
|
||||
DA3_MAX_PROCESS_RES = 1008
|
||||
@@ -34,6 +36,8 @@ IMAGE_NAME_SUFFIXES = (
|
||||
".bmp",
|
||||
".tga",
|
||||
)
|
||||
|
||||
|
||||
VIDEO_NAME_SUFFIXES = (
|
||||
".mp4",
|
||||
".mov",
|
||||
@@ -48,17 +52,12 @@ VIDEO_NAME_SUFFIXES = (
|
||||
|
||||
|
||||
def image_depth_model_items(owner, context):
|
||||
from .models import model_enum_items
|
||||
from .client import model_ready
|
||||
from .models import depth_plane_model_enum_items
|
||||
|
||||
return [
|
||||
item
|
||||
for item in model_enum_items(owner, context)
|
||||
if model_ready(item[0])
|
||||
]
|
||||
return depth_plane_model_enum_items(owner, context)
|
||||
|
||||
|
||||
_last_depth_plane_model = "BASE"
|
||||
_last_depth_plane_model = "MOGE2_VITB_NORMAL"
|
||||
|
||||
|
||||
def is_image_empty(candidate):
|
||||
@@ -800,6 +799,264 @@ def create_depth_image(name, depth_values):
|
||||
return image
|
||||
|
||||
|
||||
def create_moge_mask_image(name, mask_values):
|
||||
import numpy as np
|
||||
|
||||
height, width = mask_values.shape
|
||||
image = bpy.data.images.new(
|
||||
name,
|
||||
width=width,
|
||||
height=height,
|
||||
alpha=True,
|
||||
float_buffer=True,
|
||||
)
|
||||
image.colorspace_settings.name = "Non-Color"
|
||||
values = np.flipud(mask_values.astype(np.float32))
|
||||
rgba = np.ones((height, width, 4), dtype=np.float32)
|
||||
rgba[..., :3] = values[..., None]
|
||||
rgba[..., 3] = values
|
||||
image.pixels.foreach_set(rgba.ravel())
|
||||
image.pack()
|
||||
return image
|
||||
|
||||
|
||||
def create_moge_normal_image(name, normal_values):
|
||||
import numpy as np
|
||||
|
||||
converted = normal_values.astype(np.float32, copy=True)
|
||||
converted *= np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
||||
converted = np.clip(converted * 0.5 + 0.5, 0.0, 1.0)
|
||||
height, width = converted.shape[:2]
|
||||
image = bpy.data.images.new(
|
||||
name,
|
||||
width=width,
|
||||
height=height,
|
||||
alpha=False,
|
||||
float_buffer=True,
|
||||
)
|
||||
image.colorspace_settings.name = "Non-Color"
|
||||
rgba = np.ones((height, width, 4), dtype=np.float32)
|
||||
rgba[..., :3] = np.flipud(converted)
|
||||
image.pixels.foreach_set(rgba.ravel())
|
||||
image.pack()
|
||||
return image
|
||||
|
||||
|
||||
def create_moge_sequence_image(name, frames, scene, directory, prefix, kind):
|
||||
import numpy as np
|
||||
|
||||
directory = Path(directory)
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
height, width = frames[0].shape[:2]
|
||||
image = bpy.data.images.new(
|
||||
name,
|
||||
width=width,
|
||||
height=height,
|
||||
alpha=True,
|
||||
float_buffer=True,
|
||||
)
|
||||
image.colorspace_settings.name = "Non-Color"
|
||||
image_settings = scene.render.image_settings
|
||||
original_format = image_settings.file_format
|
||||
paths = []
|
||||
try:
|
||||
image_settings.file_format = "OPEN_EXR"
|
||||
for index, values in enumerate(frames):
|
||||
rgba = np.ones((height, width, 4), dtype=np.float32)
|
||||
if kind == "depth":
|
||||
flipped = np.flipud(values.astype(np.float32))
|
||||
rgba[..., :3] = flipped[..., None]
|
||||
elif kind == "mask":
|
||||
flipped = np.flipud(values.astype(np.float32))
|
||||
rgba[..., :3] = flipped[..., None]
|
||||
rgba[..., 3] = flipped
|
||||
else:
|
||||
converted = values.astype(np.float32, copy=True)
|
||||
converted *= np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
||||
rgba[..., :3] = np.flipud(
|
||||
np.clip(converted * 0.5 + 0.5, 0.0, 1.0)
|
||||
)
|
||||
image.pixels.foreach_set(rgba.ravel())
|
||||
path = directory / f"{prefix}.{index + 1:04d}.exr"
|
||||
image.save_render(str(path), scene=scene)
|
||||
paths.append(path)
|
||||
finally:
|
||||
image_settings.file_format = original_format
|
||||
bpy.data.images.remove(image, do_unlink=True)
|
||||
sequence = bpy.data.images.load(str(paths[0]), check_existing=False)
|
||||
force_sequence_source(sequence)
|
||||
sequence.name = name
|
||||
sequence.colorspace_settings.name = "Non-Color"
|
||||
pack_sequence_if_possible(sequence, len(paths))
|
||||
return sequence
|
||||
|
||||
|
||||
def moge_material_node_group():
|
||||
existing = bpy.data.node_groups.get(MOGE_MATERIAL_NODE_GROUP_NAME)
|
||||
if existing is not None:
|
||||
return existing
|
||||
group = bpy.data.node_groups.new(
|
||||
MOGE_MATERIAL_NODE_GROUP_NAME,
|
||||
"ShaderNodeTree",
|
||||
)
|
||||
for name, socket_type in (
|
||||
("Color", "NodeSocketColor"),
|
||||
("Alpha", "NodeSocketFloat"),
|
||||
("Depth", "NodeSocketFloat"),
|
||||
("Bump Scale", "NodeSocketFloat"),
|
||||
):
|
||||
socket = group.interface.new_socket(
|
||||
name=name,
|
||||
in_out="INPUT",
|
||||
socket_type=socket_type,
|
||||
)
|
||||
if name == "Alpha":
|
||||
socket.default_value = 1.0
|
||||
if name == "Bump Scale":
|
||||
socket.default_value = 1.0
|
||||
for name, socket_type in (
|
||||
("Color", "NodeSocketColor"),
|
||||
("Roughness", "NodeSocketFloat"),
|
||||
("Alpha", "NodeSocketFloat"),
|
||||
("Normal", "NodeSocketVector"),
|
||||
("Displacement", "NodeSocketVector"),
|
||||
):
|
||||
group.interface.new_socket(
|
||||
name=name,
|
||||
in_out="OUTPUT",
|
||||
socket_type=socket_type,
|
||||
)
|
||||
nodes = group.nodes
|
||||
links = group.links
|
||||
group_input = nodes.new("NodeGroupInput")
|
||||
group_input.location = (-900.0, 100.0)
|
||||
bump = nodes.new("ShaderNodeBump")
|
||||
bump.location = (250.0, -100.0)
|
||||
bump_distance = nodes.new("ShaderNodeMath")
|
||||
bump_distance.operation = "MULTIPLY"
|
||||
bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE
|
||||
bump_distance.location = (0.0, -100.0)
|
||||
background = nodes.new("ShaderNodeAttribute")
|
||||
background.attribute_type = "GEOMETRY"
|
||||
background.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE
|
||||
background.location = (-650.0, -400.0)
|
||||
material_mode = nodes.new("ShaderNodeAttribute")
|
||||
material_mode.attribute_type = "GEOMETRY"
|
||||
material_mode.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE
|
||||
material_mode.location = (-650.0, -600.0)
|
||||
strength = nodes.new("ShaderNodeAttribute")
|
||||
strength.attribute_type = "GEOMETRY"
|
||||
strength.attribute_name = DEPTH_SCALE_ATTRIBUTE
|
||||
strength.location = (-400.0, -600.0)
|
||||
clamp_depth = nodes.new("ShaderNodeMath")
|
||||
clamp_depth.operation = "MAXIMUM"
|
||||
clamp_depth.location = (-350.0, -350.0)
|
||||
enable_depth = nodes.new("ShaderNodeMath")
|
||||
enable_depth.operation = "MULTIPLY"
|
||||
enable_depth.location = (-100.0, -350.0)
|
||||
scale_depth = nodes.new("ShaderNodeMath")
|
||||
scale_depth.operation = "MULTIPLY"
|
||||
scale_depth.location = (150.0, -350.0)
|
||||
displacement = nodes.new("ShaderNodeDisplacement")
|
||||
displacement.inputs["Midlevel"].default_value = 0.0
|
||||
displacement.inputs["Scale"].default_value = 1.0
|
||||
displacement.location = (400.0, -350.0)
|
||||
group_output = nodes.new("NodeGroupOutput")
|
||||
group_output.location = (300.0, 100.0)
|
||||
links.new(group_input.outputs["Color"], group_output.inputs["Color"])
|
||||
links.new(group_input.outputs["Alpha"], group_output.inputs["Alpha"])
|
||||
links.new(group_input.outputs["Color"], bump.inputs["Height"])
|
||||
links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0])
|
||||
links.new(bump_distance.outputs[0], bump.inputs["Distance"])
|
||||
links.new(bump.outputs["Normal"], group_output.inputs["Normal"])
|
||||
links.new(group_input.outputs["Depth"], clamp_depth.inputs[0])
|
||||
links.new(background.outputs["Fac"], clamp_depth.inputs[1])
|
||||
links.new(clamp_depth.outputs[0], enable_depth.inputs[0])
|
||||
links.new(material_mode.outputs["Fac"], enable_depth.inputs[1])
|
||||
links.new(enable_depth.outputs[0], scale_depth.inputs[0])
|
||||
links.new(strength.outputs["Fac"], scale_depth.inputs[1])
|
||||
links.new(scale_depth.outputs[0], displacement.inputs["Height"])
|
||||
links.new(displacement.outputs["Displacement"], group_output.inputs["Displacement"])
|
||||
group_output.inputs["Roughness"].default_value = 0.5
|
||||
group.color_tag = "SHADER"
|
||||
group.description = "Apply MoGe-2 color and derive bump detail from it."
|
||||
group.use_fake_user = True
|
||||
return group
|
||||
|
||||
|
||||
def create_moge_material(
|
||||
name,
|
||||
color_image,
|
||||
depth_image,
|
||||
depth_scale,
|
||||
unshaded=False,
|
||||
frame_count=1,
|
||||
frame_start=1,
|
||||
use_cyclic=False,
|
||||
):
|
||||
material = bpy.data.materials.new(name)
|
||||
material.use_nodes = True
|
||||
if hasattr(material, "surface_render_method"):
|
||||
material.surface_render_method = "DITHERED"
|
||||
nodes = material.node_tree.nodes
|
||||
nodes.clear()
|
||||
links = material.node_tree.links
|
||||
uv = nodes.new("ShaderNodeUVMap")
|
||||
uv.uv_map = "UVMap"
|
||||
uv.location = (-900.0, 200.0)
|
||||
color = nodes.new("ShaderNodeTexImage")
|
||||
color.image = color_image
|
||||
color.location = (-650.0, 300.0)
|
||||
depth = nodes.new("ShaderNodeTexImage")
|
||||
depth.image = depth_image
|
||||
depth.interpolation = "Linear"
|
||||
depth.location = (-650.0, -100.0)
|
||||
if frame_count > 1:
|
||||
configure_image_user(
|
||||
color.image_user,
|
||||
frame_count,
|
||||
frame_start=frame_start,
|
||||
use_cyclic=use_cyclic,
|
||||
)
|
||||
configure_image_user(
|
||||
depth.image_user,
|
||||
frame_count,
|
||||
frame_start=frame_start,
|
||||
use_cyclic=use_cyclic,
|
||||
)
|
||||
layer = nodes.new("ShaderNodeGroup")
|
||||
layer.node_tree = moge_material_node_group()
|
||||
layer.name = MOGE_MATERIAL_NODE_GROUP_NAME
|
||||
layer.location = (-300.0, 200.0)
|
||||
output = nodes.new("ShaderNodeOutputMaterial")
|
||||
output.location = (500.0, 200.0)
|
||||
links.new(uv.outputs["UV"], color.inputs["Vector"])
|
||||
links.new(uv.outputs["UV"], depth.inputs["Vector"])
|
||||
links.new(color.outputs["Color"], layer.inputs["Color"])
|
||||
links.new(color.outputs["Alpha"], layer.inputs["Alpha"])
|
||||
links.new(depth.outputs["Color"], layer.inputs["Depth"])
|
||||
layer.inputs["Bump Scale"].default_value = depth_scale
|
||||
if unshaded:
|
||||
shader = nodes.new("ShaderNodeEmission")
|
||||
transparent = nodes.new("ShaderNodeBsdfTransparent")
|
||||
mix = nodes.new("ShaderNodeMixShader")
|
||||
links.new(layer.outputs["Color"], shader.inputs["Color"])
|
||||
links.new(layer.outputs["Alpha"], mix.inputs["Fac"])
|
||||
links.new(transparent.outputs["BSDF"], mix.inputs[1])
|
||||
links.new(shader.outputs["Emission"], mix.inputs[2])
|
||||
links.new(mix.outputs["Shader"], output.inputs["Surface"])
|
||||
return material
|
||||
shader = nodes.new("ShaderNodeBsdfPrincipled")
|
||||
links.new(layer.outputs["Color"], shader.inputs["Base Color"])
|
||||
links.new(layer.outputs["Alpha"], shader.inputs["Alpha"])
|
||||
links.new(layer.outputs["Roughness"], shader.inputs["Roughness"])
|
||||
links.new(layer.outputs["Normal"], shader.inputs["Normal"])
|
||||
links.new(shader.outputs["BSDF"], output.inputs["Surface"])
|
||||
links.new(layer.outputs["Displacement"], output.inputs["Displacement"])
|
||||
set_displacement_only(material)
|
||||
return material
|
||||
|
||||
|
||||
def grid_segments(bounds, subdivision_level):
|
||||
if subdivision_level < 0:
|
||||
raise ValueError("Subdivision must not be negative")
|
||||
@@ -1193,8 +1450,8 @@ def node_asset_path():
|
||||
return Path(__file__).with_name("assets") / NODE_ASSET_NAME
|
||||
|
||||
|
||||
def load_depth_plane_node_group():
|
||||
existing = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME)
|
||||
def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME):
|
||||
existing = bpy.data.node_groups.get(group_name)
|
||||
if existing is not None:
|
||||
return existing
|
||||
asset_path = node_asset_path()
|
||||
@@ -1204,26 +1461,26 @@ def load_depth_plane_node_group():
|
||||
if not use_pack:
|
||||
directory = str(asset_path) + "/NodeTree/"
|
||||
result = bpy.ops.wm.append(
|
||||
filepath=directory + DEPTH_PLANE_NODE_GROUP_NAME,
|
||||
filepath=directory + group_name,
|
||||
directory=directory,
|
||||
filename=DEPTH_PLANE_NODE_GROUP_NAME,
|
||||
filename=group_name,
|
||||
link=False,
|
||||
do_reuse_local_id=True,
|
||||
)
|
||||
if "FINISHED" not in result:
|
||||
raise RuntimeError("Unable to append the Geometry Nodes asset")
|
||||
node_group = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME)
|
||||
node_group = bpy.data.node_groups.get(group_name)
|
||||
if node_group is None:
|
||||
raise RuntimeError("Appended Geometry Nodes asset is unavailable")
|
||||
node_group.use_fake_user = True
|
||||
return node_group
|
||||
with bpy.data.libraries.load(str(asset_path), link=True) as (source, target):
|
||||
if DEPTH_PLANE_NODE_GROUP_NAME not in source.node_groups:
|
||||
if group_name not in source.node_groups:
|
||||
raise RuntimeError(
|
||||
"Geometry Nodes asset is missing "
|
||||
f"{DEPTH_PLANE_NODE_GROUP_NAME}"
|
||||
f"{group_name}"
|
||||
)
|
||||
target.node_groups = [DEPTH_PLANE_NODE_GROUP_NAME]
|
||||
target.node_groups = [group_name]
|
||||
linked_group = target.node_groups[0]
|
||||
node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group)
|
||||
if linked_group != node_group and linked_group.users == 0:
|
||||
@@ -1278,6 +1535,42 @@ def add_depth_plane_processing(
|
||||
return modifier
|
||||
|
||||
|
||||
def add_moge_depth_plane_processing(
|
||||
plane_object,
|
||||
depth_image,
|
||||
material,
|
||||
frame_start=1,
|
||||
frame_duration=1,
|
||||
use_cyclic=False,
|
||||
animate_frame=False,
|
||||
):
|
||||
node_group = load_depth_plane_node_group(
|
||||
MOGE_DEPTH_PLANE_NODE_GROUP_NAME
|
||||
)
|
||||
modifier = plane_object.modifiers.new(
|
||||
name=MOGE_DEPTH_PLANE_NODE_GROUP_NAME,
|
||||
type="NODES",
|
||||
)
|
||||
modifier.node_group = node_group
|
||||
modifier.show_group_selector = False
|
||||
for name, value in (
|
||||
("Depth Image", depth_image),
|
||||
("Material", material),
|
||||
("Frame", 1),
|
||||
("Frame Start", frame_start),
|
||||
("Frames", max(int(frame_duration or 1), 1)),
|
||||
("Cyclic", use_cyclic),
|
||||
):
|
||||
modifier[modifier_input_identifier(node_group, name)] = value
|
||||
if animate_frame:
|
||||
frame_identifier = modifier_input_identifier(node_group, "Frame")
|
||||
driver = modifier.driver_add(f'["{frame_identifier}"]')
|
||||
driver.driver.type = "SCRIPTED"
|
||||
driver.driver.expression = "frame"
|
||||
plane_object.update_tag(refresh={"DATA"})
|
||||
return modifier
|
||||
|
||||
|
||||
def finish_depth_plane(
|
||||
context,
|
||||
source_object,
|
||||
@@ -1573,6 +1866,152 @@ class BuildDepthPlane(bpy.types.Operator):
|
||||
return {"FINISHED"}
|
||||
|
||||
|
||||
def build_moge2_depth_plane(
|
||||
context,
|
||||
source_object,
|
||||
npz_path,
|
||||
delete_source,
|
||||
unshaded=False,
|
||||
color_image=None,
|
||||
foreground_image=None,
|
||||
data_frames_directory=None,
|
||||
):
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
with np.load(npz_path, allow_pickle=False) as archive:
|
||||
depth = np.asarray(archive["depth"], dtype=np.float32)
|
||||
mask = np.asarray(archive["mask"], dtype=bool)
|
||||
intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32)
|
||||
except (KeyError, OSError) as error:
|
||||
raise ValueError(f"Invalid MoGe-2 NPZ: {npz_path}") from error
|
||||
if depth.ndim == 2:
|
||||
depth = depth[None, ...]
|
||||
mask = mask[None, ...]
|
||||
if intrinsics.ndim == 2:
|
||||
intrinsics = intrinsics[None, ...]
|
||||
if not (len(depth) == len(mask) == len(intrinsics)):
|
||||
raise ValueError("MoGe-2 arrays have different frame counts")
|
||||
|
||||
first_valid = mask[0] & np.isfinite(depth[0]) & (depth[0] > 0.0)
|
||||
if not first_valid.any():
|
||||
raise ValueError("MoGe-2 mask contains no valid depth")
|
||||
reference_depth = float(np.median(depth[0][first_valid]))
|
||||
height, width = depth[0].shape
|
||||
pixel_intrinsics = intrinsics[0].copy()
|
||||
pixel_intrinsics[0] *= width
|
||||
pixel_intrinsics[1] *= height
|
||||
bounds, _ = camera_plane_bounds(
|
||||
source_object,
|
||||
depth[0],
|
||||
pixel_intrinsics,
|
||||
1.0,
|
||||
reference_depth=reference_depth,
|
||||
)
|
||||
displacement = [reference_depth - frame for frame in depth]
|
||||
source_frame_start, _, _, source_use_cyclic = image_user_settings(
|
||||
source_object
|
||||
)
|
||||
mesh = create_plane_mesh(source_object, bounds, None)
|
||||
plane_object = bpy.data.objects.new(source_object.name, mesh)
|
||||
context.collection.objects.link(plane_object)
|
||||
depth_scale = apply_plane_transform(plane_object, source_object, bounds)
|
||||
apply_object_scale_to_mesh(plane_object)
|
||||
depth_image = create_depth_image(
|
||||
f"{depth_data_name(source_object)}.moge2",
|
||||
displacement[0] * depth_scale,
|
||||
)
|
||||
if len(depth) > 1 and data_frames_directory is not None:
|
||||
bpy.data.images.remove(depth_image, do_unlink=True)
|
||||
depth_image = create_moge_sequence_image(
|
||||
f"{depth_data_name(source_object)}.moge2",
|
||||
[frame * depth_scale for frame in displacement],
|
||||
context.scene,
|
||||
data_frames_directory,
|
||||
"moge_depth",
|
||||
"depth",
|
||||
)
|
||||
material = create_moge_material(
|
||||
source_object.name,
|
||||
foreground_image or color_image or source_object.data,
|
||||
depth_image,
|
||||
depth_scale,
|
||||
unshaded=unshaded,
|
||||
frame_count=len(depth),
|
||||
frame_start=source_frame_start,
|
||||
use_cyclic=source_use_cyclic,
|
||||
)
|
||||
mesh.materials.append(material)
|
||||
modifier = add_moge_depth_plane_processing(
|
||||
plane_object,
|
||||
depth_image,
|
||||
material,
|
||||
frame_start=source_frame_start,
|
||||
frame_duration=len(depth),
|
||||
use_cyclic=source_use_cyclic,
|
||||
animate_frame=is_animated_image(source_object.data),
|
||||
)
|
||||
plane_object["snapform_reference_depth"] = reference_depth
|
||||
plane_object["snapform_depth_scale"] = depth_scale
|
||||
plane_object["snapform_frame_count"] = len(depth)
|
||||
return finish_depth_plane(
|
||||
context,
|
||||
source_object,
|
||||
plane_object,
|
||||
delete_source,
|
||||
)
|
||||
|
||||
|
||||
class BuildMoge2DepthPlane(bpy.types.Operator):
|
||||
bl_idname = "snapform.build_moge2_depth_plane"
|
||||
bl_label = "Build MoGe-2 Depth Plane"
|
||||
bl_options = {"INTERNAL"}
|
||||
|
||||
snapform_source_object_name: bpy.props.StringProperty(options={"HIDDEN"})
|
||||
snapform_npz_path: bpy.props.StringProperty(options={"HIDDEN"})
|
||||
snapform_background_path: bpy.props.StringProperty(options={"HIDDEN"})
|
||||
snapform_color_frames_dir: bpy.props.StringProperty(options={"HIDDEN"})
|
||||
snapform_data_frames_dir: bpy.props.StringProperty(options={"HIDDEN"})
|
||||
snapform_delete_source: bpy.props.BoolProperty(default=True, options={"HIDDEN"})
|
||||
snapform_unshaded: bpy.props.BoolProperty(default=False, options={"HIDDEN"})
|
||||
|
||||
def execute(self, context):
|
||||
source_object = bpy.data.objects.get(self.snapform_source_object_name)
|
||||
if not is_image_empty(source_object):
|
||||
self.report({"ERROR"}, "The source image empty is no longer available")
|
||||
return {"CANCELLED"}
|
||||
try:
|
||||
foreground_image = None
|
||||
if self.snapform_background_path:
|
||||
foreground_image = create_removed_background_image(
|
||||
source_object,
|
||||
self.snapform_background_path,
|
||||
assign_to_source=False,
|
||||
)
|
||||
color_image = None
|
||||
if self.snapform_color_frames_dir:
|
||||
color_image = create_sequence_image(
|
||||
f"{image_base_name(source_object.data)}.frame",
|
||||
self.snapform_color_frames_dir,
|
||||
)
|
||||
build_moge2_depth_plane(
|
||||
context,
|
||||
source_object,
|
||||
Path(self.snapform_npz_path),
|
||||
self.snapform_delete_source,
|
||||
unshaded=self.snapform_unshaded,
|
||||
color_image=color_image,
|
||||
foreground_image=foreground_image,
|
||||
data_frames_directory=(
|
||||
self.snapform_data_frames_dir or None
|
||||
),
|
||||
)
|
||||
except (OSError, RuntimeError, ValueError) as error:
|
||||
self.report({"ERROR"}, str(error))
|
||||
return {"CANCELLED"}
|
||||
return {"FINISHED"}
|
||||
|
||||
|
||||
class BuildOriginalDepthPlane(bpy.types.Operator):
|
||||
bl_idname = "snapform.build_original_depth_plane"
|
||||
bl_label = "Build Original Depth Plane"
|
||||
@@ -1746,19 +2185,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
bl_description = "Convert an image to a displaced depth plane"
|
||||
bl_options = {"REGISTER"}
|
||||
|
||||
snapform_depth_source: bpy.props.EnumProperty(
|
||||
name="Depth from",
|
||||
description="Source used to displace the plane",
|
||||
items=(
|
||||
(
|
||||
"ESTIMATE",
|
||||
"AI Generated",
|
||||
"Generate depth with an AI model",
|
||||
),
|
||||
("ORIGINAL", "Original", "Use the original image as depth"),
|
||||
),
|
||||
default="ORIGINAL",
|
||||
)
|
||||
snapform_unshaded: bpy.props.BoolProperty(
|
||||
name="Unshaded",
|
||||
description="Display the source color through an Emission shader",
|
||||
@@ -1771,8 +2197,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
)
|
||||
snapform_model: bpy.props.EnumProperty(
|
||||
name="Model",
|
||||
description="Model used to estimate image depth",
|
||||
description="Model or original image used to build depth",
|
||||
items=image_depth_model_items,
|
||||
default=1,
|
||||
)
|
||||
snapform_use_ray_pose: bpy.props.BoolProperty(
|
||||
name="Use Ray Pose",
|
||||
@@ -1780,6 +2207,13 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
default=True,
|
||||
)
|
||||
snapform_process_res_factor: process_resolution_scale_property()
|
||||
snapform_moge_resolution_level: bpy.props.IntProperty(
|
||||
name="Resolution Level",
|
||||
description="MoGe-2 inference resolution level; higher is finer and slower",
|
||||
default=5,
|
||||
min=0,
|
||||
max=9,
|
||||
)
|
||||
snapform_keep_original: bpy.props.BoolProperty(
|
||||
name="Keep Original",
|
||||
description="Keep the source Image Empty after conversion",
|
||||
@@ -1789,14 +2223,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
@classmethod
|
||||
def poll(cls, context):
|
||||
from .operators import active_job
|
||||
from .client import service_busy, model_ready
|
||||
from .models import MODELS
|
||||
from .client import service_busy
|
||||
|
||||
return (
|
||||
is_image_empty(context.object)
|
||||
and active_job() is None
|
||||
and not service_busy()
|
||||
and any(model_ready(model_key) for model_key in MODELS)
|
||||
)
|
||||
|
||||
def invoke(self, context, _event):
|
||||
@@ -1816,17 +2248,19 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
|
||||
def draw(self, context):
|
||||
layout = self.layout
|
||||
layout.prop(self, "snapform_depth_source")
|
||||
if self.snapform_depth_source == "ESTIMATE":
|
||||
layout.prop(self, "snapform_model")
|
||||
if self.snapform_model in {"SMALL", "BASE"}:
|
||||
estimate_box = layout.box()
|
||||
available_models = image_depth_model_items(self, context)
|
||||
if not available_models:
|
||||
estimate_box.label(text="No DA3 models installed", icon="ERROR")
|
||||
estimate_box.label(text="Download one in Models")
|
||||
else:
|
||||
estimate_box.prop(self, "snapform_model")
|
||||
estimate_box.prop(self, "snapform_use_ray_pose")
|
||||
draw_process_resolution(estimate_box, self, context.object.data)
|
||||
estimate_box.prop(self, "snapform_use_ray_pose")
|
||||
draw_process_resolution(estimate_box, self, context.object.data)
|
||||
elif self.snapform_model.startswith("MOGE2_"):
|
||||
from .client import model_ready
|
||||
|
||||
moge_box = layout.box()
|
||||
if not model_ready(self.snapform_model):
|
||||
moge_box.label(text="MoGe-2 model is not installed", icon="ERROR")
|
||||
moge_box.label(text="Download it in Models")
|
||||
moge_box.prop(self, "snapform_moge_resolution_level")
|
||||
layout.prop(self, "snapform_remove_background")
|
||||
layout.prop(self, "snapform_unshaded")
|
||||
layout.separator(type="LINE")
|
||||
@@ -1835,6 +2269,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
def execute(self, context):
|
||||
from .operators import (
|
||||
IMAGE_DEPTH_JOB,
|
||||
MOGE2_DEPTH_JOB,
|
||||
ORIGINAL_BACKGROUND_JOB,
|
||||
ORIGINAL_FRAMES_JOB,
|
||||
)
|
||||
@@ -1844,7 +2279,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
context.scene.snapform_settings.video_frames = (
|
||||
default_extract_frames(source_object)
|
||||
)
|
||||
if self.snapform_depth_source == "ORIGINAL":
|
||||
if self.snapform_model == "ORIGINAL":
|
||||
try:
|
||||
source_path, delete_input = prepare_image_input(
|
||||
source_object.data,
|
||||
@@ -1871,6 +2306,34 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
remove_temporary_input(source_path)
|
||||
return {"CANCELLED"}
|
||||
|
||||
if self.snapform_model.startswith("MOGE2_"):
|
||||
try:
|
||||
source_path, delete_input = prepare_image_input(
|
||||
source_object.data,
|
||||
context.scene,
|
||||
)
|
||||
except RuntimeError as error:
|
||||
self.report({"ERROR"}, str(error))
|
||||
return {"CANCELLED"}
|
||||
result = bpy.ops.snapform.run_job(
|
||||
job_type=MOGE2_DEPTH_JOB,
|
||||
snapform_input_path=str(source_path),
|
||||
snapform_source_object_name=source_object.name,
|
||||
snapform_moge_resolution_level=(
|
||||
self.snapform_moge_resolution_level
|
||||
),
|
||||
snapform_model_key=self.snapform_model,
|
||||
snapform_delete_source=not self.snapform_keep_original,
|
||||
snapform_delete_input=delete_input,
|
||||
snapform_unshaded=self.snapform_unshaded,
|
||||
snapform_remove_background=self.snapform_remove_background,
|
||||
)
|
||||
if "CANCELLED" not in result:
|
||||
return {"FINISHED"}
|
||||
if delete_input:
|
||||
remove_temporary_input(source_path)
|
||||
return {"CANCELLED"}
|
||||
|
||||
image = source_object.data
|
||||
try:
|
||||
process_res = scaled_process_resolution(
|
||||
|
||||
Reference in New Issue
Block a user