nan
c5b77c5c85
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing. - Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing. - Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Developed unit tests for server functionalities in `test_server.py`, including model caching and job management. - Enhanced job server to handle frame extraction without requiring models.
2212 lines
76 KiB
Python
2212 lines
76 KiB
Python
import re
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import bpy
|
|
|
|
|
|
DEPTH_KEY_CANDIDATES = ("depth", "depths")
|
|
INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts")
|
|
DISPLACEMENT_STRENGTH = 1.0
|
|
DISPLACEMENT_MIDLEVEL = 0.0
|
|
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"
|
|
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 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
|
|
IMAGE_NAME_SUFFIXES = (
|
|
".jpeg",
|
|
".jpg",
|
|
".png",
|
|
".webp",
|
|
".tiff",
|
|
".tif",
|
|
".exr",
|
|
".hdr",
|
|
".bmp",
|
|
".tga",
|
|
)
|
|
|
|
|
|
VIDEO_NAME_SUFFIXES = (
|
|
".mp4",
|
|
".mov",
|
|
".avi",
|
|
".mkv",
|
|
".webm",
|
|
".m4v",
|
|
".mts",
|
|
".mpg",
|
|
".mpeg",
|
|
)
|
|
|
|
|
|
def image_depth_model_items(owner, context):
|
|
from ..properties import depth_plane_model_enum_items
|
|
|
|
return depth_plane_model_enum_items(owner, context)
|
|
|
|
|
|
def is_image_empty(candidate):
|
|
return bool(
|
|
candidate
|
|
and getattr(candidate, "type", None) == "EMPTY"
|
|
and getattr(candidate, "empty_display_type", None) == "IMAGE"
|
|
and getattr(candidate, "data", None) is not None
|
|
)
|
|
|
|
|
|
def is_animated_image(image):
|
|
return getattr(image, "source", None) in {"MOVIE", "SEQUENCE"}
|
|
|
|
|
|
def video_total_frames(image):
|
|
"""Return the total animated frame count, or 1 for a still image."""
|
|
source = getattr(image, "source", None)
|
|
if source == "MOVIE":
|
|
return int(getattr(image, "frame_duration", 1) or 1)
|
|
if source == "SEQUENCE":
|
|
return len(collect_sequence_frames(image)) or 1
|
|
return 1
|
|
|
|
|
|
def image_user_settings(source_object):
|
|
"""Return the Empty's image user settings as a four-item tuple."""
|
|
image_user = getattr(source_object, "image_user", None)
|
|
if image_user is None:
|
|
return 1, 0, 0, False
|
|
return (
|
|
int(getattr(image_user, "frame_start", 1) or 1),
|
|
int(getattr(image_user, "frame_offset", 0) or 0),
|
|
int(getattr(image_user, "frame_duration", 0) or 0),
|
|
bool(getattr(image_user, "use_cyclic", False)),
|
|
)
|
|
|
|
|
|
def extract_frame_range(source_object):
|
|
"""Return the clamped (start_frame, frame_count) range extracted from the source."""
|
|
total_frames = video_total_frames(getattr(source_object, "data", None))
|
|
if total_frames <= 1:
|
|
return 1, total_frames
|
|
_frame_start, frame_offset, image_frames, _use_cyclic = (
|
|
image_user_settings(source_object)
|
|
)
|
|
frame_count = (
|
|
min(image_frames, total_frames)
|
|
if image_frames >= 1
|
|
else total_frames
|
|
)
|
|
start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
|
|
start_frame = max(start_frame, 1)
|
|
return start_frame, frame_count
|
|
|
|
|
|
def default_extract_frames(source_object):
|
|
"""Return the Empty's Image Frames value, or the full frame count when unset."""
|
|
return extract_frame_range(source_object)[1]
|
|
|
|
|
|
def image_path(image):
|
|
filepath_from_user = getattr(image, "filepath_from_user", None)
|
|
filepath = (
|
|
filepath_from_user()
|
|
if callable(filepath_from_user)
|
|
else getattr(image, "filepath", "")
|
|
)
|
|
if not filepath:
|
|
return Path()
|
|
return Path(bpy.path.abspath(filepath)).expanduser()
|
|
|
|
|
|
def image_is_packed(image):
|
|
if getattr(image, "packed_file", None) is not None:
|
|
return True
|
|
return bool(getattr(image, "packed_files", ()))
|
|
|
|
|
|
def remove_temporary_input(path):
|
|
"""Remove a temporary input file or frame directory."""
|
|
path = Path(path)
|
|
try:
|
|
if path.is_dir():
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
else:
|
|
path.unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def image_base_name(image):
|
|
name = image.name
|
|
lower_name = name.lower()
|
|
for suffix in IMAGE_NAME_SUFFIXES + VIDEO_NAME_SUFFIXES:
|
|
suffix_index = lower_name.rfind(suffix)
|
|
if suffix_index <= 0:
|
|
continue
|
|
trailing = name[suffix_index + len(suffix) :]
|
|
if trailing and not (
|
|
trailing.startswith(".") and trailing[1:].isdigit()
|
|
):
|
|
continue
|
|
return f"{name[:suffix_index]}{trailing}"
|
|
return name
|
|
|
|
|
|
def depth_data_name(source_object):
|
|
return f"{image_base_name(source_object.data)}.depth"
|
|
|
|
|
|
def normal_data_name(source_object):
|
|
return f"{image_base_name(source_object.data)}.normal"
|
|
|
|
|
|
def mask_data_name(source_object):
|
|
return f"{image_base_name(source_object.data)}.mask"
|
|
|
|
|
|
def copy_depth_image(source_object):
|
|
"""Return a renamed copy of the source image for use as the depth map."""
|
|
depth_image = source_object.data.copy()
|
|
depth_image.name = depth_data_name(source_object)
|
|
return depth_image
|
|
|
|
|
|
def removed_background_data_name(source_object):
|
|
return f"{image_base_name(source_object.data)}.foreground"
|
|
|
|
|
|
def scaled_process_resolution(image, scale_percent):
|
|
"""Return the scaled long edge clamped and aligned for DA3 inference."""
|
|
width, height = image.size
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError("The source image has no readable dimensions")
|
|
if scale_percent < 1.0 or scale_percent > 100.0:
|
|
raise ValueError("Resolution Scale must be between 1% and 100%")
|
|
process_res = max(width, height) * scale_percent / 100.0
|
|
process_res = max(DA3_MIN_PROCESS_RES, process_res)
|
|
process_res = min(DA3_MAX_PROCESS_RES, process_res)
|
|
return int(process_res / DA3_PATCH_SIZE + 0.5) * DA3_PATCH_SIZE
|
|
|
|
|
|
def process_resolution_scale_property():
|
|
return bpy.props.FloatProperty(
|
|
name="Resolution Scale",
|
|
description=(
|
|
"Scale the source image long edge for DA3 processing; the result "
|
|
"is limited to 252-1008 px and aligned to a multiple of 14"
|
|
),
|
|
default=100.0,
|
|
min=1.0,
|
|
max=100.0,
|
|
step=100,
|
|
precision=0,
|
|
subtype="PERCENTAGE",
|
|
)
|
|
|
|
|
|
def draw_process_resolution(layout, owner, image):
|
|
process_res = scaled_process_resolution(
|
|
image,
|
|
owner.process_res_factor,
|
|
)
|
|
layout.prop(
|
|
owner,
|
|
"process_res_factor",
|
|
text=f"Resolution Scale ({process_res} px)",
|
|
)
|
|
|
|
|
|
def save_packed_image(image, scene):
|
|
output_path = Path(tempfile.mkdtemp(prefix="snapform-input-")) / "input.png"
|
|
image_settings = scene.render.image_settings
|
|
original_file_format = image_settings.file_format
|
|
try:
|
|
image_settings.file_format = "PNG"
|
|
image.save_render(str(output_path), scene=scene)
|
|
except (OSError, RuntimeError) as error:
|
|
output_path.unlink(missing_ok=True)
|
|
raise RuntimeError("Unable to export the packed image") from error
|
|
finally:
|
|
image_settings.file_format = original_file_format
|
|
if not output_path.is_file():
|
|
raise RuntimeError("Blender did not export the packed image")
|
|
return output_path
|
|
|
|
|
|
def collect_sequence_frames(image):
|
|
"""Return the numbered frame files of a SEQUENCE image, or [] when absent."""
|
|
source_path = image_path(image)
|
|
if not is_animated_image(image) or not source_path.is_file():
|
|
return []
|
|
match = re.match(
|
|
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
|
|
source_path.name,
|
|
)
|
|
if match is None:
|
|
return []
|
|
prefix = match.group("prefix")
|
|
suffix = match.group("suffix").lower()
|
|
frames = []
|
|
for path in source_path.parent.iterdir():
|
|
frame_match = re.match(
|
|
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
|
|
path.name,
|
|
)
|
|
if (
|
|
frame_match is not None
|
|
and frame_match.group("prefix") == prefix
|
|
and frame_match.group("suffix").lower() == suffix
|
|
):
|
|
frames.append((int(frame_match.group("digits")), path))
|
|
return [path for _digits, path in sorted(frames)]
|
|
|
|
|
|
def prepare_image_input(image, scene):
|
|
source_path = image_path(image)
|
|
if image_is_packed(image):
|
|
return save_packed_image(image, scene), True
|
|
if getattr(image, "source", None) == "MOVIE":
|
|
if not source_path.is_file():
|
|
raise RuntimeError(
|
|
"The image empty must reference an existing video file"
|
|
)
|
|
return source_path, False
|
|
sequence_frames = collect_sequence_frames(image)
|
|
if sequence_frames:
|
|
temporary_directory = Path(tempfile.mkdtemp(prefix="snapform-input-"))
|
|
try:
|
|
for index, frame_path in enumerate(sequence_frames):
|
|
suffix = frame_path.suffix.lower() or ".jpg"
|
|
shutil.copy2(
|
|
frame_path,
|
|
temporary_directory / f"frame.{index + 1:04d}{suffix}",
|
|
)
|
|
except Exception:
|
|
shutil.rmtree(temporary_directory, ignore_errors=True)
|
|
raise
|
|
return temporary_directory, True
|
|
if not source_path.is_file():
|
|
raise RuntimeError(
|
|
"The image empty must reference an existing image file"
|
|
)
|
|
return source_path, False
|
|
|
|
|
|
def configure_image_user(
|
|
image_user,
|
|
frame_count=None,
|
|
frame_start=1,
|
|
frame_offset=0,
|
|
use_cyclic=None,
|
|
):
|
|
"""Map the image to scene frames and enable auto refresh."""
|
|
if image_user is None:
|
|
return
|
|
if hasattr(image_user, "frame_start"):
|
|
image_user.frame_start = frame_start
|
|
if hasattr(image_user, "frame_offset"):
|
|
image_user.frame_offset = frame_offset
|
|
if hasattr(image_user, "use_auto_refresh"):
|
|
image_user.use_auto_refresh = True
|
|
if use_cyclic is not None and hasattr(image_user, "use_cyclic"):
|
|
image_user.use_cyclic = use_cyclic
|
|
if frame_count is not None and hasattr(image_user, "frame_duration"):
|
|
image_user.frame_duration = frame_count
|
|
|
|
|
|
def pack_sequence_if_possible(image, expected_frames):
|
|
"""Pack a frame sequence when Blender supports it; otherwise keep external files."""
|
|
if expected_frames <= 0:
|
|
return False
|
|
try:
|
|
image.pack()
|
|
except (OSError, RuntimeError):
|
|
return False
|
|
if len(getattr(image, "packed_files", ())) >= expected_frames:
|
|
return True
|
|
try:
|
|
image.unpack(method="REMOVE")
|
|
except (OSError, RuntimeError):
|
|
pass
|
|
return False
|
|
|
|
|
|
def sorted_frame_paths(directory, prefix=None):
|
|
"""Return image files in a directory sorted by their trailing frame number."""
|
|
def frame_key(path):
|
|
match = re.search(r"(\d+)(?=\.[^.]+$)", path.name)
|
|
return (int(match.group(1)) if match is not None else 0, path.name)
|
|
|
|
return sorted(
|
|
(
|
|
path
|
|
for path in Path(directory).iterdir()
|
|
if path.is_file()
|
|
and path.suffix.lower() in IMAGE_NAME_SUFFIXES
|
|
and (prefix is None or path.stem.startswith(prefix))
|
|
),
|
|
key=frame_key,
|
|
)
|
|
|
|
|
|
def force_sequence_source(image):
|
|
"""Ensure a loaded image is a SEQUENCE datablock for its numbered frames."""
|
|
if getattr(image, "source", None) == "SEQUENCE":
|
|
return image
|
|
try:
|
|
image.source = "SEQUENCE"
|
|
except (OSError, RuntimeError, TypeError) as error:
|
|
raise RuntimeError(
|
|
"Blender cannot load the frame sequence"
|
|
) from error
|
|
if getattr(image, "source", None) != "SEQUENCE":
|
|
raise RuntimeError("Blender cannot load the frame sequence")
|
|
return image
|
|
|
|
|
|
def create_sequence_image(name, frame_directory, frame_prefix="frame."):
|
|
"""Load a numbered frame directory as a SEQUENCE image datablock."""
|
|
frame_paths = sorted_frame_paths(frame_directory, frame_prefix)
|
|
if not frame_paths:
|
|
raise RuntimeError(f"No image frames found in: {frame_directory}")
|
|
try:
|
|
image = bpy.data.images.load(
|
|
str(frame_paths[0]),
|
|
check_existing=False,
|
|
)
|
|
force_sequence_source(image)
|
|
except (OSError, RuntimeError) as error:
|
|
if "image" in locals():
|
|
bpy.data.images.remove(image, do_unlink=True)
|
|
raise RuntimeError("Unable to load the frame sequence") from error
|
|
image.name = name
|
|
pack_sequence_if_possible(image, len(frame_paths))
|
|
return image
|
|
|
|
|
|
def create_depth_sequence_image(name, depth_frames, scene, frame_directory):
|
|
"""Write one EXR per depth frame and return a SEQUENCE image datablock."""
|
|
import numpy as np
|
|
|
|
frame_directory = Path(frame_directory)
|
|
frame_directory.mkdir(parents=True, exist_ok=True)
|
|
height, width = depth_frames[0].shape
|
|
image = bpy.data.images.new(
|
|
name,
|
|
width=width,
|
|
height=height,
|
|
alpha=False,
|
|
float_buffer=True,
|
|
)
|
|
image.colorspace_settings.name = "Non-Color"
|
|
image_settings = scene.render.image_settings
|
|
original_file_format = image_settings.file_format
|
|
frame_paths = []
|
|
try:
|
|
image_settings.file_format = "OPEN_EXR"
|
|
for frame_index, depth_values in enumerate(depth_frames):
|
|
# NPZ arrays use a top-left origin, while Blender image pixels start below.
|
|
depth_pixels = np.flipud(depth_values)
|
|
rgba = np.empty((height, width, 4), dtype=np.float32)
|
|
rgba[..., 0] = depth_pixels
|
|
rgba[..., 1] = depth_pixels
|
|
rgba[..., 2] = depth_pixels
|
|
rgba[..., 3] = 1.0
|
|
image.pixels.foreach_set(rgba.ravel())
|
|
frame_path = (
|
|
frame_directory / f"depth.{frame_index + 1:04d}.exr"
|
|
)
|
|
image.save_render(str(frame_path), scene=scene)
|
|
if not frame_path.is_file():
|
|
raise RuntimeError(f"Blender did not write: {frame_path}")
|
|
frame_paths.append(frame_path)
|
|
finally:
|
|
image_settings.file_format = original_file_format
|
|
bpy.data.images.remove(image, do_unlink=True)
|
|
if not frame_paths:
|
|
raise RuntimeError("Unable to write the depth frame sequence")
|
|
try:
|
|
sequence_image = bpy.data.images.load(
|
|
str(frame_paths[0]),
|
|
check_existing=False,
|
|
)
|
|
force_sequence_source(sequence_image)
|
|
except (OSError, RuntimeError) as error:
|
|
if "sequence_image" in locals():
|
|
bpy.data.images.remove(sequence_image, do_unlink=True)
|
|
raise RuntimeError(
|
|
"Unable to load the depth frame sequence"
|
|
) from error
|
|
sequence_image.name = name
|
|
sequence_image.colorspace_settings.name = "Non-Color"
|
|
pack_sequence_if_possible(sequence_image, len(frame_paths))
|
|
return sequence_image
|
|
|
|
|
|
def create_removed_background_image(
|
|
source_object,
|
|
output_path,
|
|
assign_to_source=True,
|
|
):
|
|
output_path = Path(output_path)
|
|
if output_path.is_dir():
|
|
foreground_paths = sorted_frame_paths(output_path, "foreground.")
|
|
if not foreground_paths:
|
|
raise RuntimeError(
|
|
f"Background removal result has no foreground frames: "
|
|
f"{output_path}"
|
|
)
|
|
try:
|
|
result_image = bpy.data.images.load(
|
|
str(foreground_paths[0]),
|
|
check_existing=False,
|
|
)
|
|
result_image.alpha_mode = "CHANNEL_PACKED"
|
|
force_sequence_source(result_image)
|
|
except (OSError, RuntimeError) as error:
|
|
if "result_image" in locals():
|
|
bpy.data.images.remove(result_image, do_unlink=True)
|
|
raise RuntimeError("Unable to load the background result") from error
|
|
source_image = source_object.data
|
|
original_fake_user = source_image.use_fake_user
|
|
try:
|
|
result_image.name = removed_background_data_name(source_object)
|
|
pack_sequence_if_possible(result_image, len(foreground_paths))
|
|
if assign_to_source:
|
|
image_user = getattr(source_object, "image_user", None)
|
|
(
|
|
original_frame_start,
|
|
_original_frame_offset,
|
|
_original_frame_duration,
|
|
original_use_cyclic,
|
|
) = image_user_settings(source_object)
|
|
source_image.use_fake_user = True
|
|
source_object.data = result_image
|
|
configure_image_user(
|
|
image_user,
|
|
len(foreground_paths),
|
|
frame_start=original_frame_start,
|
|
use_cyclic=original_use_cyclic,
|
|
)
|
|
except (OSError, RuntimeError) as error:
|
|
source_image.use_fake_user = original_fake_user
|
|
bpy.data.images.remove(result_image)
|
|
raise RuntimeError("Unable to load the background result") from error
|
|
return result_image
|
|
|
|
if not output_path.is_file():
|
|
raise RuntimeError(
|
|
f"Background removal result does not exist: {output_path}"
|
|
)
|
|
try:
|
|
result_image = bpy.data.images.load(
|
|
str(output_path),
|
|
check_existing=False,
|
|
)
|
|
result_image.alpha_mode = "CHANNEL_PACKED"
|
|
except RuntimeError as error:
|
|
raise RuntimeError("Unable to load the background result") from error
|
|
source_image = source_object.data
|
|
original_fake_user = source_image.use_fake_user
|
|
try:
|
|
result_image.name = removed_background_data_name(source_object)
|
|
result_image.pack()
|
|
if assign_to_source:
|
|
source_image.use_fake_user = True
|
|
source_object.data = result_image
|
|
except (OSError, RuntimeError) as error:
|
|
source_image.use_fake_user = original_fake_user
|
|
bpy.data.images.remove(result_image)
|
|
raise RuntimeError("Unable to pack the background result") from error
|
|
return result_image
|
|
|
|
|
|
def plane_bounds(source_object):
|
|
image_width, image_height = source_object.data.size
|
|
if image_width <= 0 or image_height <= 0:
|
|
raise ValueError("The source image has no readable dimensions")
|
|
|
|
aspect = image_width / image_height
|
|
display_size = float(source_object.empty_display_size)
|
|
if aspect >= 1.0:
|
|
width = display_size
|
|
height = display_size / aspect
|
|
else:
|
|
width = display_size * aspect
|
|
height = display_size
|
|
|
|
offset = source_object.empty_image_offset
|
|
x_min = float(offset[0]) * width
|
|
y_min = float(offset[1]) * height
|
|
return (
|
|
x_min,
|
|
x_min + width,
|
|
y_min,
|
|
y_min + height,
|
|
)
|
|
|
|
|
|
def plane_fit_transform(source_object, camera_bounds):
|
|
target_bounds = plane_bounds(source_object)
|
|
camera_width = camera_bounds[1] - camera_bounds[0]
|
|
camera_height = camera_bounds[3] - camera_bounds[2]
|
|
if camera_width <= 0.0 or camera_height <= 0.0:
|
|
raise ValueError("The camera plane has invalid dimensions")
|
|
|
|
target_width = target_bounds[1] - target_bounds[0]
|
|
target_height = target_bounds[3] - target_bounds[2]
|
|
scale_x = target_width / camera_width
|
|
scale_y = target_height / camera_height
|
|
scale_z = (scale_x * scale_y) ** 0.5
|
|
camera_center_x = (camera_bounds[0] + camera_bounds[1]) * 0.5
|
|
camera_center_y = (camera_bounds[2] + camera_bounds[3]) * 0.5
|
|
target_center_x = (target_bounds[0] + target_bounds[1]) * 0.5
|
|
target_center_y = (target_bounds[2] + target_bounds[3]) * 0.5
|
|
return (
|
|
scale_x,
|
|
scale_y,
|
|
scale_z,
|
|
target_center_x - camera_center_x * scale_x,
|
|
target_center_y - camera_center_y * scale_y,
|
|
)
|
|
|
|
|
|
def apply_plane_transform(plane_object, source_object, camera_bounds):
|
|
from mathutils import Matrix
|
|
|
|
scale_x, scale_y, scale_z, offset_x, offset_y = plane_fit_transform(
|
|
source_object,
|
|
camera_bounds,
|
|
)
|
|
fit_matrix = Matrix(
|
|
(
|
|
(scale_x, 0.0, 0.0, offset_x),
|
|
(0.0, scale_y, 0.0, offset_y),
|
|
(0.0, 0.0, 1.0, 0.0),
|
|
(0.0, 0.0, 0.0, 1.0),
|
|
)
|
|
)
|
|
plane_object.matrix_world = source_object.matrix_world @ fit_matrix
|
|
return abs(float(plane_object.scale.z)) * abs(scale_z)
|
|
|
|
|
|
def apply_object_scale_to_mesh(plane_object):
|
|
"""Bake object scale into the plane while keeping location and rotation."""
|
|
from mathutils import Matrix
|
|
|
|
scale = plane_object.scale.copy()
|
|
plane_object.data.transform(
|
|
Matrix.Diagonal((scale.x, scale.y, scale.z, 1.0))
|
|
)
|
|
plane_object.data.update()
|
|
plane_object.scale = (1.0, 1.0, 1.0)
|
|
|
|
|
|
def scale_depth_values(depth, depth_scale):
|
|
return prepare_depth_values(depth) * depth_scale
|
|
|
|
|
|
def extract_depth_arrays(npz_path):
|
|
try:
|
|
import numpy as np
|
|
except ImportError as error:
|
|
raise RuntimeError("Blender Python does not provide NumPy") from error
|
|
|
|
try:
|
|
with np.load(npz_path, allow_pickle=False) as archive:
|
|
depth_key = next(
|
|
(
|
|
key
|
|
for key in DEPTH_KEY_CANDIDATES
|
|
if key in archive.files
|
|
),
|
|
None,
|
|
)
|
|
if depth_key is None:
|
|
keys = ", ".join(archive.files) or "none"
|
|
raise ValueError(
|
|
f"Mini NPZ has no depth array; available keys: {keys}"
|
|
)
|
|
depth = np.asarray(archive[depth_key], dtype=np.float32)
|
|
except OSError as error:
|
|
raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error
|
|
|
|
if depth.ndim >= 3 and depth.shape[-1] == 1:
|
|
depth = depth[..., 0]
|
|
if depth.ndim == 2:
|
|
depth = depth[None, ...]
|
|
if depth.ndim != 3:
|
|
raise ValueError(
|
|
f"Unsupported depth array shape: {tuple(depth.shape)}"
|
|
)
|
|
return [depth[index] for index in range(depth.shape[0])]
|
|
|
|
|
|
def extract_depth_array(npz_path):
|
|
return extract_depth_arrays(npz_path)[0]
|
|
|
|
|
|
def extract_intrinsics_array(npz_path):
|
|
try:
|
|
import numpy as np
|
|
except ImportError as error:
|
|
raise RuntimeError("Blender Python does not provide NumPy") from error
|
|
|
|
try:
|
|
with np.load(npz_path, allow_pickle=False) as archive:
|
|
intrinsics_key = next(
|
|
(
|
|
key
|
|
for key in INTRINSICS_KEY_CANDIDATES
|
|
if key in archive.files
|
|
),
|
|
None,
|
|
)
|
|
if intrinsics_key is None:
|
|
keys = ", ".join(archive.files) or "none"
|
|
raise ValueError(
|
|
"Mini NPZ has no camera intrinsics; "
|
|
f"available keys: {keys}"
|
|
)
|
|
intrinsics = np.asarray(
|
|
archive[intrinsics_key],
|
|
dtype=np.float32,
|
|
)
|
|
except OSError as error:
|
|
raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error
|
|
|
|
if intrinsics.ndim == 3:
|
|
intrinsics = intrinsics[0]
|
|
if intrinsics.shape != (3, 3):
|
|
raise ValueError(
|
|
f"Unsupported intrinsics shape: {tuple(intrinsics.shape)}"
|
|
)
|
|
return intrinsics
|
|
|
|
|
|
def prepare_depth_values(depth):
|
|
import numpy as np
|
|
|
|
valid = np.isfinite(depth)
|
|
if not valid.any():
|
|
raise ValueError("Depth array contains no finite values")
|
|
values = depth.astype(np.float32, copy=True)
|
|
values[~valid] = 0.0
|
|
return values
|
|
|
|
|
|
def depth_displacement_values(depth):
|
|
"""Convert raw DA3 camera distance to positive outward displacement."""
|
|
import numpy as np
|
|
|
|
return 1.0 - np.asarray(depth, dtype=np.float32)
|
|
|
|
|
|
def camera_plane_bounds(
|
|
source_object,
|
|
depth,
|
|
intrinsics,
|
|
displacement_strength,
|
|
reference_depth=None,
|
|
):
|
|
import numpy as np
|
|
|
|
scale = abs(float(displacement_strength))
|
|
if scale <= 1e-8:
|
|
raise ValueError("Displacement Strength must not be zero")
|
|
|
|
if reference_depth is None:
|
|
valid = np.isfinite(depth) & (depth > 0.0)
|
|
if not valid.any():
|
|
raise ValueError("Depth array contains no positive finite values")
|
|
reference_depth = float(np.median(depth[valid]))
|
|
reference_distance = reference_depth * scale
|
|
|
|
height, width = depth.shape
|
|
fx = float(intrinsics[0, 0])
|
|
fy = float(intrinsics[1, 1])
|
|
cx = float(intrinsics[0, 2])
|
|
cy = float(intrinsics[1, 2])
|
|
if fx <= 0.0 or fy <= 0.0:
|
|
raise ValueError("Camera intrinsics contain an invalid focal length")
|
|
|
|
source_width, source_height = source_object.data.size
|
|
if source_width <= 0 or source_height <= 0:
|
|
raise ValueError("The source image has no readable dimensions")
|
|
source_aspect = source_width / source_height
|
|
|
|
# Predicted fx/fy may differ even for square pixels. Preserve their projected
|
|
# area while enforcing the source image aspect ratio to avoid image stretch.
|
|
raw_width = width * reference_distance / fx
|
|
raw_height = height * reference_distance / fy
|
|
plane_area = raw_width * raw_height
|
|
plane_width = float(np.sqrt(plane_area * source_aspect))
|
|
plane_height = float(np.sqrt(plane_area / source_aspect))
|
|
|
|
x_min = -(cx / width) * plane_width
|
|
x_max = (1.0 - cx / width) * plane_width
|
|
y_min = -(1.0 - cy / height) * plane_height
|
|
y_max = (cy / height) * plane_height
|
|
|
|
# Preserve an intentional Image Empty origin offset around the camera frame.
|
|
offset = source_object.empty_image_offset
|
|
x_offset = (float(offset[0]) + 0.5) * plane_width
|
|
y_offset = (float(offset[1]) + 0.5) * plane_height
|
|
return (
|
|
(
|
|
x_min + x_offset,
|
|
x_max + x_offset,
|
|
y_min + y_offset,
|
|
y_max + y_offset,
|
|
),
|
|
reference_depth,
|
|
)
|
|
|
|
|
|
def create_depth_image(name, depth_values):
|
|
import numpy as np
|
|
|
|
height, width = depth_values.shape
|
|
image = bpy.data.images.new(
|
|
name,
|
|
width=width,
|
|
height=height,
|
|
alpha=False,
|
|
float_buffer=True,
|
|
)
|
|
image.colorspace_settings.name = "Non-Color"
|
|
image.file_format = "OPEN_EXR"
|
|
|
|
# NPZ arrays use a top-left origin, while Blender image pixels start below.
|
|
depth_pixels = np.flipud(depth_values)
|
|
rgba = np.empty((height, width, 4), dtype=np.float32)
|
|
rgba[..., 0] = depth_pixels
|
|
rgba[..., 1] = depth_pixels
|
|
rgba[..., 2] = depth_pixels
|
|
rgba[..., 3] = 1.0
|
|
image.pixels.foreach_set(rgba.ravel())
|
|
# Pack as EXR so values outside 0..1 survive saving and reopening the blend.
|
|
image.pack()
|
|
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, convert_coordinates=True):
|
|
import numpy as np
|
|
|
|
converted = normal_values.astype(np.float32, copy=True)
|
|
if convert_coordinates:
|
|
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.interpolation = "Linear"
|
|
color.extension = "EXTEND"
|
|
color.location = (-650.0, 300.0)
|
|
depth = nodes.new("ShaderNodeTexImage")
|
|
depth.image = depth_image
|
|
depth.interpolation = "Linear"
|
|
depth.extension = "EXTEND"
|
|
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")
|
|
|
|
width = bounds[1] - bounds[0]
|
|
height = bounds[3] - bounds[2]
|
|
if width <= 0.0 or height <= 0.0:
|
|
raise ValueError("The plane has invalid dimensions")
|
|
|
|
short_segments = 2**subdivision_level
|
|
if width >= height:
|
|
return (
|
|
max(short_segments, int(short_segments * width / height + 0.5)),
|
|
short_segments,
|
|
)
|
|
return (
|
|
short_segments,
|
|
max(short_segments, int(short_segments * height / width + 0.5)),
|
|
)
|
|
|
|
|
|
def create_plane_mesh(source_object, bounds=None, subdivision_level=5):
|
|
x_min, x_max, y_min, y_max = (
|
|
bounds if bounds is not None else plane_bounds(source_object)
|
|
)
|
|
if subdivision_level is None:
|
|
x_segments, y_segments = 1, 1
|
|
else:
|
|
x_segments, y_segments = grid_segments(
|
|
(x_min, x_max, y_min, y_max),
|
|
subdivision_level,
|
|
)
|
|
vertices = []
|
|
uv_coordinates = []
|
|
for y_index in range(y_segments + 1):
|
|
v = y_index / y_segments
|
|
y = y_min + (y_max - y_min) * v
|
|
for x_index in range(x_segments + 1):
|
|
u = x_index / x_segments
|
|
x = x_min + (x_max - x_min) * u
|
|
vertices.append((x, y, 0.0))
|
|
uv_coordinates.append((u, v))
|
|
|
|
faces = []
|
|
row_width = x_segments + 1
|
|
for y_index in range(y_segments):
|
|
for x_index in range(x_segments):
|
|
bottom_left = y_index * row_width + x_index
|
|
faces.append(
|
|
(
|
|
bottom_left,
|
|
bottom_left + 1,
|
|
bottom_left + row_width + 1,
|
|
bottom_left + row_width,
|
|
)
|
|
)
|
|
|
|
mesh = bpy.data.meshes.new(source_object.name)
|
|
mesh.from_pydata(vertices, (), faces)
|
|
mesh.update()
|
|
for polygon in mesh.polygons:
|
|
polygon.use_smooth = True
|
|
uv_layer = mesh.uv_layers.new(name="UVMap")
|
|
for loop in mesh.loops:
|
|
uv_layer.data[loop.index].uv = uv_coordinates[
|
|
loop.vertex_index
|
|
]
|
|
return mesh
|
|
|
|
|
|
def material_node_group(depth_plane=False, surface_group_name=None):
|
|
group_name = (
|
|
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME
|
|
if depth_plane
|
|
else surface_group_name or SURFACE_MATERIAL_NODE_GROUP_NAME
|
|
)
|
|
existing = bpy.data.node_groups.get(group_name)
|
|
if existing is not None:
|
|
return existing
|
|
group = bpy.data.node_groups.new(group_name, "ShaderNodeTree")
|
|
interface = group.interface
|
|
input_sockets = [
|
|
("Color", "NodeSocketColor"),
|
|
("Alpha", "NodeSocketFloat"),
|
|
("Bump Scale", "NodeSocketFloat"),
|
|
]
|
|
if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME:
|
|
input_sockets.insert(2, ("Normal", "NodeSocketColor"))
|
|
input_sockets.insert(3, ("Normal Strength", "NodeSocketFloat"))
|
|
if depth_plane:
|
|
input_sockets.insert(2, ("Depth", "NodeSocketFloat"))
|
|
for name, socket_type in input_sockets:
|
|
socket = interface.new_socket(
|
|
name=name,
|
|
in_out="INPUT",
|
|
socket_type=socket_type,
|
|
)
|
|
if name == "Alpha":
|
|
socket.default_value = 1.0
|
|
socket.min_value = 0.0
|
|
socket.max_value = 1.0
|
|
socket.subtype = "FACTOR"
|
|
elif name == "Bump Scale":
|
|
socket.default_value = 1.0
|
|
socket.min_value = -10000.0
|
|
socket.max_value = 10000.0
|
|
elif name == "Normal Strength":
|
|
socket.default_value = 1.0
|
|
socket.min_value = 0.0
|
|
socket.max_value = 10.0
|
|
output_sockets = [
|
|
("Color", "NodeSocketColor"),
|
|
("Roughness", "NodeSocketFloat"),
|
|
("Alpha", "NodeSocketFloat"),
|
|
("Normal", "NodeSocketVector"),
|
|
]
|
|
if depth_plane:
|
|
output_sockets.append(("Displacement", "NodeSocketVector"))
|
|
for name, socket_type in output_sockets:
|
|
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 = (-500.0, 100.0)
|
|
group_output = nodes.new("NodeGroupOutput")
|
|
group_output.location = (300.0, 100.0)
|
|
roughness_curve = nodes.new("ShaderNodeFloatCurve")
|
|
roughness_curve.mapping.initialize()
|
|
curve_points = roughness_curve.mapping.curves[0].points
|
|
curve_points[0].location = (0.0, 0.5)
|
|
curve_points[-1].location = (1.0, 0.1)
|
|
roughness_curve.mapping.update()
|
|
roughness_curve.location = (-250.0, 350.0)
|
|
bump = nodes.new("ShaderNodeBump")
|
|
bump.location = (-250.0, 0.0)
|
|
if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME:
|
|
normal_map = nodes.new("ShaderNodeNormalMap")
|
|
normal_map.space = "OBJECT"
|
|
normal_map.location = (-500.0, -150.0)
|
|
links.new(group_input.outputs["Normal"], normal_map.inputs["Color"])
|
|
links.new(
|
|
group_input.outputs["Normal Strength"],
|
|
normal_map.inputs["Strength"],
|
|
)
|
|
links.new(normal_map.outputs["Normal"], bump.inputs["Normal"])
|
|
displacement = None
|
|
|
|
processed_alpha = group_input.outputs["Alpha"]
|
|
bump_depth = group_input.outputs["Color"]
|
|
if depth_plane:
|
|
displacement = nodes.new("ShaderNodeDisplacement")
|
|
displacement.inputs["Midlevel"].default_value = 0.0
|
|
displacement.inputs["Scale"].default_value = 1.0
|
|
group_input.location = (-1200.0, 100.0)
|
|
group_output.location = (1050.0, 100.0)
|
|
roughness_curve.location = (-300.0, 450.0)
|
|
bump.location = (750.0, 50.0)
|
|
displacement.location = (750.0, -300.0)
|
|
background_attribute = nodes.new("ShaderNodeAttribute")
|
|
background_attribute.attribute_type = "GEOMETRY"
|
|
background_attribute.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE
|
|
background_attribute.location = (-950.0, -500.0)
|
|
material_attribute = nodes.new("ShaderNodeAttribute")
|
|
material_attribute.attribute_type = "GEOMETRY"
|
|
material_attribute.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE
|
|
material_attribute.location = (-50.0, -650.0)
|
|
strength_attribute = nodes.new("ShaderNodeAttribute")
|
|
strength_attribute.attribute_type = "GEOMETRY"
|
|
strength_attribute.attribute_name = DEPTH_SCALE_ATTRIBUTE
|
|
strength_attribute.location = (200.0, -650.0)
|
|
foreground_mask = nodes.new("ShaderNodeMath")
|
|
foreground_mask.operation = "GREATER_THAN"
|
|
foreground_mask.inputs[1].default_value = 0.5
|
|
foreground_mask.location = (-950.0, -250.0)
|
|
clamp_foreground = nodes.new("ShaderNodeMath")
|
|
clamp_foreground.operation = "MAXIMUM"
|
|
clamp_foreground.location = (-650.0, -400.0)
|
|
masked_depth = nodes.new("ShaderNodeMix")
|
|
masked_depth.data_type = "FLOAT"
|
|
masked_depth.location = (-350.0, -350.0)
|
|
enabled_depth = nodes.new("ShaderNodeMath")
|
|
enabled_depth.operation = "MULTIPLY"
|
|
enabled_depth.location = (250.0, -350.0)
|
|
scaled_depth = nodes.new("ShaderNodeMath")
|
|
scaled_depth.operation = "MULTIPLY"
|
|
scaled_depth.location = (450.0, -350.0)
|
|
links.new(group_input.outputs["Alpha"], foreground_mask.inputs[0])
|
|
links.new(group_input.outputs["Depth"], clamp_foreground.inputs[0])
|
|
links.new(background_attribute.outputs["Fac"], clamp_foreground.inputs[1])
|
|
links.new(foreground_mask.outputs[0], masked_depth.inputs[0])
|
|
links.new(background_attribute.outputs["Fac"], masked_depth.inputs[2])
|
|
links.new(clamp_foreground.outputs[0], masked_depth.inputs[3])
|
|
links.new(masked_depth.outputs["Result"], enabled_depth.inputs[0])
|
|
links.new(material_attribute.outputs["Fac"], enabled_depth.inputs[1])
|
|
links.new(enabled_depth.outputs[0], scaled_depth.inputs[0])
|
|
links.new(strength_attribute.outputs["Fac"], scaled_depth.inputs[1])
|
|
processed_alpha = foreground_mask.outputs[0]
|
|
processed_depth = scaled_depth.outputs[0]
|
|
|
|
links.new(group_input.outputs["Color"], group_output.inputs["Color"])
|
|
links.new(group_input.outputs["Color"], roughness_curve.inputs["Value"])
|
|
links.new(
|
|
roughness_curve.outputs["Value"],
|
|
group_output.inputs["Roughness"],
|
|
)
|
|
links.new(processed_alpha, group_output.inputs["Alpha"])
|
|
links.new(bump_depth, bump.inputs["Height"])
|
|
bump_distance = nodes.new("ShaderNodeMath")
|
|
bump_distance.operation = "MULTIPLY"
|
|
bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE
|
|
bump_distance.location = (500.0, 100.0) if depth_plane else (-250.0, 150.0)
|
|
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"])
|
|
if displacement is not None:
|
|
links.new(processed_depth, displacement.inputs["Height"])
|
|
links.new(
|
|
displacement.outputs["Displacement"],
|
|
group_output.inputs["Displacement"],
|
|
)
|
|
group.color_tag = "SHADER"
|
|
group.description = (
|
|
"Prepare DA3 depth-plane material outputs."
|
|
if depth_plane
|
|
else f"Prepare {group_name} material outputs."
|
|
)
|
|
group.use_fake_user = True
|
|
return group
|
|
|
|
|
|
def create_material(
|
|
name,
|
|
color_image,
|
|
depth_scale,
|
|
displacement_image=None,
|
|
foreground_image=None,
|
|
displacement_midlevel=0.0,
|
|
displacement_scale=1.0,
|
|
unshaded=False,
|
|
frame_count=1,
|
|
frame_start=1,
|
|
color_frame_offset=0,
|
|
color_frame_duration=None,
|
|
foreground_frame_offset=0,
|
|
foreground_frame_duration=None,
|
|
depth_frame_offset=0,
|
|
depth_frame_duration=None,
|
|
use_cyclic=None,
|
|
surface_group_name=None,
|
|
normal_image=None,
|
|
):
|
|
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()
|
|
|
|
color_texture = nodes.new("ShaderNodeTexImage")
|
|
uv_map = nodes.new("ShaderNodeUVMap")
|
|
output = nodes.new("ShaderNodeOutputMaterial")
|
|
|
|
uv_map.uv_map = "UVMap"
|
|
uv_map.location = (-1250.0, 350.0)
|
|
material_image = foreground_image or color_image
|
|
color_texture.image = material_image
|
|
color_texture.interpolation = "Linear"
|
|
color_texture.extension = "EXTEND"
|
|
color_texture.location = (-1000.0, 350.0)
|
|
output.location = (850.0, 350.0)
|
|
extracted_frames = frame_count if frame_count > 1 else None
|
|
color_frames = (
|
|
color_frame_duration
|
|
if color_frame_duration is not None
|
|
else extracted_frames
|
|
)
|
|
depth_frames = (
|
|
depth_frame_duration
|
|
if depth_frame_duration is not None
|
|
else extracted_frames
|
|
)
|
|
material_frames = color_frames
|
|
material_frame_offset = color_frame_offset
|
|
if foreground_image is not None:
|
|
material_frames = (
|
|
foreground_frame_duration
|
|
if foreground_frame_duration is not None
|
|
else extracted_frames
|
|
)
|
|
material_frame_offset = foreground_frame_offset
|
|
if frame_count > 1 or is_animated_image(material_image):
|
|
configure_image_user(
|
|
color_texture.image_user,
|
|
material_frames,
|
|
frame_start=frame_start,
|
|
frame_offset=material_frame_offset,
|
|
use_cyclic=use_cyclic,
|
|
)
|
|
|
|
links = material.node_tree.links
|
|
links.new(uv_map.outputs["UV"], color_texture.inputs["Vector"])
|
|
alpha = color_texture.outputs["Alpha"]
|
|
material_inputs = nodes.new("ShaderNodeGroup")
|
|
is_depth_plane = displacement_image is not None
|
|
material_inputs.node_tree = material_node_group(
|
|
is_depth_plane,
|
|
surface_group_name=surface_group_name,
|
|
)
|
|
material_inputs.name = material_inputs.node_tree.name
|
|
material_inputs.label = material_inputs.name
|
|
material_inputs.location = (-300.0, 350.0)
|
|
material_inputs.inputs["Bump Scale"].default_value = depth_scale
|
|
links.new(color_texture.outputs["Color"], material_inputs.inputs["Color"])
|
|
links.new(alpha, material_inputs.inputs["Alpha"])
|
|
if normal_image is not None and "Normal" in material_inputs.inputs:
|
|
normal_texture = nodes.new("ShaderNodeTexImage")
|
|
normal_texture.image = normal_image
|
|
normal_texture.interpolation = "Linear"
|
|
normal_texture.extension = "EXTEND"
|
|
normal_texture.location = (-1000.0, -50.0)
|
|
links.new(uv_map.outputs["UV"], normal_texture.inputs["Vector"])
|
|
links.new(normal_texture.outputs["Color"], material_inputs.inputs["Normal"])
|
|
if unshaded:
|
|
emission = nodes.new("ShaderNodeEmission")
|
|
transparent = nodes.new("ShaderNodeBsdfTransparent")
|
|
mix_shader = nodes.new("ShaderNodeMixShader")
|
|
emission.location = (0.0, 550.0)
|
|
transparent.location = (0.0, 150.0)
|
|
mix_shader.location = (300.0, 350.0)
|
|
links.new(
|
|
material_inputs.outputs["Color"],
|
|
emission.inputs["Color"],
|
|
)
|
|
links.new(
|
|
material_inputs.outputs["Alpha"],
|
|
mix_shader.inputs["Fac"],
|
|
)
|
|
links.new(
|
|
transparent.outputs["BSDF"],
|
|
mix_shader.inputs[1],
|
|
)
|
|
links.new(
|
|
emission.outputs["Emission"],
|
|
mix_shader.inputs[2],
|
|
)
|
|
links.new(
|
|
mix_shader.outputs["Shader"],
|
|
output.inputs["Surface"],
|
|
)
|
|
else:
|
|
shader = nodes.new("ShaderNodeBsdfPrincipled")
|
|
shader.location = (250.0, 400.0)
|
|
|
|
links.new(
|
|
material_inputs.outputs["Color"],
|
|
shader.inputs["Base Color"],
|
|
)
|
|
links.new(material_inputs.outputs["Alpha"], shader.inputs["Alpha"])
|
|
links.new(material_inputs.outputs["Normal"], shader.inputs["Normal"])
|
|
links.new(
|
|
material_inputs.outputs["Roughness"],
|
|
shader.inputs["Roughness"],
|
|
)
|
|
links.new(shader.outputs["BSDF"], output.inputs["Surface"])
|
|
|
|
if displacement_image is not None:
|
|
depth_texture = nodes.new("ShaderNodeTexImage")
|
|
depth_texture.image = displacement_image
|
|
depth_texture.interpolation = "Linear"
|
|
depth_texture.extension = "EXTEND"
|
|
if frame_count > 1 or is_animated_image(displacement_image):
|
|
configure_image_user(
|
|
depth_texture.image_user,
|
|
depth_frames,
|
|
frame_start=frame_start,
|
|
frame_offset=depth_frame_offset,
|
|
use_cyclic=use_cyclic,
|
|
)
|
|
depth_texture.location = (-1000.0, -400.0)
|
|
links.new(uv_map.outputs["UV"], depth_texture.inputs["Vector"])
|
|
displacement_height = depth_texture.outputs["Color"]
|
|
if displacement_midlevel != 0.0 or displacement_scale != 1.0:
|
|
depth_adjust = nodes.new("ShaderNodeMath")
|
|
depth_adjust.operation = "SUBTRACT"
|
|
depth_adjust.inputs[1].default_value = displacement_midlevel
|
|
depth_adjust.location = (500.0, -550.0)
|
|
depth_scale_node = nodes.new("ShaderNodeMath")
|
|
depth_scale_node.operation = "MULTIPLY"
|
|
depth_scale_node.inputs[1].default_value = displacement_scale
|
|
depth_scale_node.location = (700.0, -550.0)
|
|
links.new(displacement_height, depth_adjust.inputs[0])
|
|
links.new(depth_adjust.outputs[0], depth_scale_node.inputs[0])
|
|
links.new(depth_scale_node.outputs[0], material_inputs.inputs["Depth"])
|
|
else:
|
|
links.new(displacement_height, material_inputs.inputs["Depth"])
|
|
links.new(
|
|
material_inputs.outputs["Displacement"],
|
|
output.inputs["Displacement"],
|
|
)
|
|
set_displacement_only(material)
|
|
return material
|
|
|
|
|
|
def set_displacement_only(material):
|
|
for owner, property_name in (
|
|
(material, "surface_displacement_method"),
|
|
(material, "displacement_method"),
|
|
(getattr(material, "cycles", None), "displacement_method"),
|
|
):
|
|
if owner is not None and hasattr(owner, property_name):
|
|
setattr(owner, property_name, "DISPLACEMENT")
|
|
return
|
|
|
|
|
|
def node_asset_path():
|
|
return Path(__file__).parent.parent / "assets" / NODE_ASSET_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()
|
|
if not asset_path.is_file():
|
|
raise RuntimeError(f"Missing Geometry Nodes asset: {asset_path}")
|
|
use_pack = bpy.app.version >= (5, 0, 0)
|
|
if not use_pack:
|
|
directory = str(asset_path) + "/NodeTree/"
|
|
result = bpy.ops.wm.append(
|
|
filepath=directory + group_name,
|
|
directory=directory,
|
|
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(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 group_name not in source.node_groups:
|
|
raise RuntimeError(
|
|
"Geometry Nodes asset is missing "
|
|
f"{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:
|
|
bpy.data.node_groups.remove(linked_group)
|
|
node_group.use_fake_user = True
|
|
return node_group
|
|
|
|
|
|
def modifier_input_identifier(node_group, name):
|
|
for item in node_group.interface.items_tree:
|
|
if item.item_type != "SOCKET" or item.in_out != "INPUT":
|
|
continue
|
|
if item.name == name:
|
|
return item.identifier
|
|
raise RuntimeError(f"Geometry Nodes input is missing: {name}")
|
|
|
|
|
|
def modifier_input_slots(modifier):
|
|
"""Return the runtime input collection, or None for the legacy API."""
|
|
properties = getattr(modifier, "properties", None)
|
|
if properties is None:
|
|
return None
|
|
return getattr(properties, "inputs", None)
|
|
|
|
|
|
def set_modifier_input(modifier, identifier, value):
|
|
"""Assign a value to a Geometry Nodes modifier input socket."""
|
|
inputs = modifier_input_slots(modifier)
|
|
if inputs is not None:
|
|
# Blender 5.2+ exposes inputs as runtime RNA keyed by socket identifier.
|
|
getattr(inputs, identifier).value = value
|
|
else:
|
|
# Blender 4.x exposes inputs as modifier ID properties.
|
|
modifier[identifier] = value
|
|
|
|
|
|
def get_modifier_input(modifier, identifier):
|
|
"""Read a Geometry Nodes modifier input socket value."""
|
|
inputs = modifier_input_slots(modifier)
|
|
if inputs is not None:
|
|
return getattr(inputs, identifier).value
|
|
return modifier[identifier]
|
|
|
|
|
|
def modifier_input_driver_path(modifier, identifier):
|
|
"""Return the driver data path for a modifier input socket."""
|
|
if modifier_input_slots(modifier) is not None:
|
|
return (
|
|
f'modifiers["{modifier.name}"]'
|
|
f".properties.inputs.{identifier}.value"
|
|
)
|
|
return f'modifiers["{modifier.name}"]["{identifier}"]'
|
|
|
|
|
|
def drive_modifier_input(owner_object, modifier, identifier, expression="frame"):
|
|
"""Add a scripted driver to a modifier input so it follows the scene frame."""
|
|
driver = owner_object.driver_add(
|
|
modifier_input_driver_path(modifier, identifier)
|
|
)
|
|
driver.driver.type = "SCRIPTED"
|
|
driver.driver.expression = expression
|
|
return driver
|
|
|
|
|
|
def add_depth_plane_processing(
|
|
plane_object,
|
|
depth_image,
|
|
foreground_image,
|
|
material,
|
|
frame_start=1,
|
|
frame_duration=1,
|
|
use_cyclic=False,
|
|
animate_frame=False,
|
|
):
|
|
node_group = load_depth_plane_node_group()
|
|
modifier = plane_object.modifiers.new(
|
|
name=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),
|
|
("Foreground Image", foreground_image),
|
|
("Material", material),
|
|
("Frame", 1),
|
|
("Frame Start", frame_start),
|
|
("Frames", max(int(frame_duration or 1), 1)),
|
|
("Cyclic", use_cyclic),
|
|
):
|
|
identifier = modifier_input_identifier(node_group, name)
|
|
set_modifier_input(modifier, identifier, value)
|
|
if animate_frame:
|
|
frame_identifier = modifier_input_identifier(node_group, "Frame")
|
|
drive_modifier_input(plane_object, modifier, frame_identifier)
|
|
plane_object.update_tag(refresh={"DATA"})
|
|
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),
|
|
):
|
|
identifier = modifier_input_identifier(node_group, name)
|
|
set_modifier_input(modifier, identifier, value)
|
|
if animate_frame:
|
|
frame_identifier = modifier_input_identifier(node_group, "Frame")
|
|
drive_modifier_input(plane_object, modifier, frame_identifier)
|
|
plane_object.update_tag(refresh={"DATA"})
|
|
return modifier
|
|
|
|
|
|
def finish_depth_plane(
|
|
context,
|
|
source_object,
|
|
plane_object,
|
|
delete_source,
|
|
):
|
|
source_name = source_object.name
|
|
for selected in context.selected_objects:
|
|
selected.select_set(False)
|
|
plane_object.select_set(True)
|
|
context.view_layer.objects.active = plane_object
|
|
if delete_source:
|
|
bpy.data.objects.remove(source_object, do_unlink=True)
|
|
plane_object.name = source_name
|
|
return plane_object
|
|
|
|
|
|
def build_image_depth_plane(
|
|
context,
|
|
source_object,
|
|
delete_source,
|
|
unshaded=False,
|
|
color_image=None,
|
|
foreground_image=None,
|
|
depth_frames_directory=None,
|
|
depth_image=None,
|
|
):
|
|
source_name = source_object.name
|
|
frame_count = 1
|
|
frames_image = None
|
|
if depth_frames_directory is not None:
|
|
frames_image = create_sequence_image(
|
|
depth_data_name(source_object),
|
|
depth_frames_directory,
|
|
)
|
|
frame_count = len(
|
|
sorted_frame_paths(depth_frames_directory, "frame.")
|
|
)
|
|
(
|
|
source_frame_start,
|
|
source_frame_offset,
|
|
source_frame_duration,
|
|
source_use_cyclic,
|
|
) = image_user_settings(source_object)
|
|
uses_extracted = frames_image is not None
|
|
texture_frame_count = (
|
|
frame_count if uses_extracted else source_frame_duration
|
|
)
|
|
texture_frame_offset = 0 if uses_extracted else source_frame_offset
|
|
depth_source_image = depth_image or source_object.data
|
|
if depth_image is None:
|
|
if frames_image is not None:
|
|
depth_source_image = frames_image
|
|
elif not is_animated_image(source_object.data):
|
|
depth_source_image = copy_depth_image(source_object)
|
|
foreground_image = (
|
|
foreground_image or color_image or source_object.data
|
|
)
|
|
mesh = create_plane_mesh(source_object, subdivision_level=None)
|
|
plane_object = bpy.data.objects.new(source_name, mesh)
|
|
context.collection.objects.link(plane_object)
|
|
plane_object.matrix_world = source_object.matrix_world.copy()
|
|
|
|
depth_scale = abs(float(plane_object.scale.z))
|
|
material = create_material(
|
|
source_name,
|
|
color_image if color_image is not None else source_object.data,
|
|
depth_scale,
|
|
displacement_image=depth_source_image,
|
|
foreground_image=foreground_image,
|
|
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
|
|
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
|
|
unshaded=unshaded,
|
|
frame_count=frame_count,
|
|
frame_start=source_frame_start,
|
|
color_frame_offset=texture_frame_offset,
|
|
color_frame_duration=texture_frame_count,
|
|
foreground_frame_offset=0,
|
|
foreground_frame_duration=frame_count,
|
|
depth_frame_offset=texture_frame_offset,
|
|
depth_frame_duration=texture_frame_count,
|
|
use_cyclic=source_use_cyclic,
|
|
)
|
|
mesh.materials.append(material)
|
|
add_depth_plane_processing(
|
|
plane_object,
|
|
depth_source_image,
|
|
foreground_image,
|
|
material,
|
|
frame_start=source_frame_start,
|
|
frame_duration=frame_count,
|
|
use_cyclic=source_use_cyclic,
|
|
animate_frame=is_animated_image(source_object.data),
|
|
)
|
|
plane_object["snapform_frame_count"] = frame_count
|
|
return finish_depth_plane(
|
|
context,
|
|
source_object,
|
|
plane_object,
|
|
delete_source,
|
|
)
|
|
|
|
|
|
def build_depth_plane(
|
|
context,
|
|
source_object,
|
|
npz_path,
|
|
delete_source,
|
|
unshaded=False,
|
|
depth_name="",
|
|
color_image=None,
|
|
foreground_image=None,
|
|
depth_frames_directory=None,
|
|
):
|
|
source_name = source_object.name
|
|
raw_depth_arrays = extract_depth_arrays(npz_path)
|
|
raw_depth = raw_depth_arrays[0]
|
|
intrinsics = extract_intrinsics_array(npz_path)
|
|
raw_depth_values = prepare_depth_values(raw_depth)
|
|
bounds, reference_depth = camera_plane_bounds(
|
|
source_object,
|
|
raw_depth_values,
|
|
intrinsics,
|
|
DISPLACEMENT_STRENGTH,
|
|
)
|
|
depth_arrays = [
|
|
depth_displacement_values(frame) for frame in raw_depth_arrays
|
|
]
|
|
depth_name = depth_name or depth_data_name(source_object)
|
|
frame_count = len(depth_arrays)
|
|
foreground_image = (
|
|
foreground_image or color_image or source_object.data
|
|
)
|
|
(
|
|
source_frame_start,
|
|
source_frame_offset,
|
|
source_frame_duration,
|
|
source_use_cyclic,
|
|
) = image_user_settings(source_object)
|
|
mesh = create_plane_mesh(source_object, bounds, None)
|
|
plane_object = bpy.data.objects.new(
|
|
source_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)
|
|
scaled_depth_arrays = [
|
|
scale_depth_values(frame, depth_scale) for frame in depth_arrays
|
|
]
|
|
depth_image = create_depth_image(
|
|
depth_name,
|
|
scaled_depth_arrays[0],
|
|
)
|
|
if frame_count > 1 and depth_frames_directory is not None:
|
|
bpy.data.images.remove(depth_image, do_unlink=True)
|
|
depth_image = create_depth_sequence_image(
|
|
depth_name,
|
|
scaled_depth_arrays,
|
|
context.scene,
|
|
depth_frames_directory,
|
|
)
|
|
|
|
material = create_material(
|
|
source_name,
|
|
color_image if color_image is not None else source_object.data,
|
|
depth_scale,
|
|
displacement_image=depth_image,
|
|
foreground_image=foreground_image,
|
|
displacement_midlevel=DISPLACEMENT_MIDLEVEL,
|
|
displacement_scale=DISPLACEMENT_STRENGTH,
|
|
unshaded=unshaded,
|
|
frame_count=frame_count,
|
|
frame_start=source_frame_start,
|
|
color_frame_offset=0 if color_image is not None else source_frame_offset,
|
|
color_frame_duration=(
|
|
frame_count if color_image is not None else source_frame_duration
|
|
),
|
|
foreground_frame_offset=0,
|
|
foreground_frame_duration=frame_count,
|
|
depth_frame_offset=0,
|
|
depth_frame_duration=frame_count,
|
|
use_cyclic=source_use_cyclic,
|
|
)
|
|
mesh.materials.append(material)
|
|
add_depth_plane_processing(
|
|
plane_object,
|
|
depth_image,
|
|
foreground_image,
|
|
material,
|
|
frame_start=source_frame_start,
|
|
frame_duration=frame_count,
|
|
use_cyclic=source_use_cyclic,
|
|
animate_frame=is_animated_image(source_object.data),
|
|
)
|
|
plane_object["snapform_reference_depth"] = reference_depth
|
|
plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0]
|
|
plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2]
|
|
plane_object["snapform_depth_scale"] = depth_scale
|
|
plane_object["snapform_frame_count"] = frame_count
|
|
|
|
return finish_depth_plane(
|
|
context,
|
|
source_object,
|
|
plane_object,
|
|
delete_source,
|
|
)
|
|
|
|
|
|
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(
|
|
depth_data_name(source_object),
|
|
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(
|
|
depth_data_name(source_object),
|
|
[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 RemoveImageBackground(bpy.types.Operator):
|
|
bl_idname = "snapform.remove_image_background"
|
|
bl_label = "Remove Background"
|
|
bl_description = "Remove the background from this Image Empty with BEN2"
|
|
bl_options = {"REGISTER"}
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
from ..job_runtime import runtime
|
|
from ..properties import BEN2_MODEL_KEY, model_ready
|
|
|
|
return (
|
|
is_image_empty(context.object)
|
|
and not runtime.server_busy()
|
|
and model_ready(BEN2_MODEL_KEY)
|
|
)
|
|
|
|
def execute(self, context):
|
|
source_object = context.object
|
|
total_frames = video_total_frames(source_object.data)
|
|
if total_frames > 1:
|
|
context.scene.snapform_settings.video_frames = (
|
|
default_extract_frames(source_object)
|
|
)
|
|
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.process_image_background(
|
|
input_path=str(source_path),
|
|
source_object_name=source_object.name,
|
|
delete_input=delete_input,
|
|
)
|
|
if "CANCELLED" not in result:
|
|
return {"FINISHED"}
|
|
if delete_input:
|
|
remove_temporary_input(source_path)
|
|
return {"CANCELLED"}
|
|
|
|
|
|
class ConvertToDepthPlane(bpy.types.Operator):
|
|
bl_idname = "snapform.convert_to_depth_plane"
|
|
bl_label = "Convert to Depth Plane"
|
|
bl_description = "Convert an image to a displaced depth plane"
|
|
bl_options = {"REGISTER"}
|
|
|
|
unshaded: bpy.props.BoolProperty(
|
|
name="Unshaded",
|
|
description="Display the source color through an Emission shader",
|
|
default=False,
|
|
)
|
|
remove_background: bpy.props.BoolProperty(
|
|
name="Remove Background",
|
|
description="Remove the background during conversion with BEN2",
|
|
default=False,
|
|
)
|
|
model: bpy.props.EnumProperty(
|
|
name="Depth From",
|
|
description="Model or original image used to build depth",
|
|
items=image_depth_model_items,
|
|
default=0,
|
|
)
|
|
use_ray_pose: bpy.props.BoolProperty(
|
|
name="Use Ray Pose",
|
|
description="Legacy task compatibility; ONNX uses the camera decoder",
|
|
default=False,
|
|
)
|
|
process_res_factor: process_resolution_scale_property()
|
|
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,
|
|
)
|
|
keep_original: bpy.props.BoolProperty(
|
|
name="Keep Original",
|
|
description="Keep the source Image Empty after conversion",
|
|
default=False,
|
|
)
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
from ..job_runtime import runtime
|
|
|
|
return (
|
|
is_image_empty(context.object)
|
|
and not runtime.server_busy()
|
|
)
|
|
|
|
def invoke(self, context, _event):
|
|
return context.window_manager.invoke_props_dialog(self, width=340)
|
|
|
|
def draw(self, context):
|
|
layout = self.layout
|
|
layout.prop(self, "model")
|
|
if self.model in {"SMALL", "BASE"}:
|
|
draw_process_resolution(layout, self, context.object.data)
|
|
elif self.model.startswith("MOGE2_"):
|
|
layout.prop(self, "resolution_level")
|
|
layout.prop(self, "remove_background")
|
|
layout.prop(self, "unshaded")
|
|
layout.separator(type="LINE")
|
|
layout.prop(self, "keep_original")
|
|
|
|
def execute(self, context):
|
|
source_object = context.object
|
|
if video_total_frames(source_object.data) > 1:
|
|
context.scene.snapform_settings.video_frames = (
|
|
default_extract_frames(source_object)
|
|
)
|
|
if self.model == "ORIGINAL":
|
|
try:
|
|
source_path, delete_input = prepare_image_input(
|
|
source_object.data,
|
|
context.scene,
|
|
)
|
|
except RuntimeError as error:
|
|
self.report({"ERROR"}, str(error))
|
|
return {"CANCELLED"}
|
|
run_original = (
|
|
bpy.ops.snapform.generate_original_background
|
|
if self.remove_background
|
|
else bpy.ops.snapform.extract_original_frames
|
|
)
|
|
result = run_original(
|
|
input_path=str(source_path),
|
|
source_object_name=source_object.name,
|
|
delete_source=not self.keep_original,
|
|
delete_input=delete_input,
|
|
unshaded=self.unshaded,
|
|
)
|
|
if "CANCELLED" not in result:
|
|
return {"FINISHED"}
|
|
if delete_input:
|
|
remove_temporary_input(source_path)
|
|
return {"CANCELLED"}
|
|
|
|
if self.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.generate_moge2_depth(
|
|
input_path=str(source_path),
|
|
source_object_name=source_object.name,
|
|
resolution_level=self.resolution_level,
|
|
model=self.model,
|
|
delete_source=not self.keep_original,
|
|
delete_input=delete_input,
|
|
unshaded=self.unshaded,
|
|
remove_background=self.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(
|
|
image,
|
|
self.process_res_factor,
|
|
)
|
|
except ValueError as error:
|
|
self.report({"ERROR"}, str(error))
|
|
return {"CANCELLED"}
|
|
try:
|
|
source_path, delete_input = prepare_image_input(
|
|
image,
|
|
context.scene,
|
|
)
|
|
except RuntimeError as error:
|
|
self.report({"ERROR"}, str(error))
|
|
return {"CANCELLED"}
|
|
|
|
result = bpy.ops.snapform.generate_da3_depth(
|
|
input_path=str(source_path),
|
|
source_object_name=source_object.name,
|
|
depth_data_name=depth_data_name(source_object),
|
|
model=self.model,
|
|
use_ray_pose=self.use_ray_pose,
|
|
process_res=process_res,
|
|
delete_source=not self.keep_original,
|
|
delete_input=delete_input,
|
|
unshaded=self.unshaded,
|
|
remove_background=self.remove_background,
|
|
)
|
|
if "CANCELLED" in result:
|
|
if delete_input:
|
|
remove_temporary_input(source_path)
|
|
return {"CANCELLED"}
|
|
return {"FINISHED"}
|
|
|
|
|
|
class SnapFormImageMenu(bpy.types.Menu):
|
|
bl_idname = "SNAPFORM_MT_image_context"
|
|
bl_label = "SnapForm"
|
|
|
|
def draw(self, _context):
|
|
from ..job_runtime import runtime
|
|
|
|
layout = self.layout
|
|
layout.operator_context = "INVOKE_DEFAULT"
|
|
if not runtime.environment_ready():
|
|
layout.operator(
|
|
"snapform.install_environment",
|
|
text="Install Environment",
|
|
icon="IMPORT",
|
|
)
|
|
return
|
|
layout.operator(
|
|
"snapform.convert_to_depth_plane",
|
|
text="Convert to Depth Plane",
|
|
icon="MOD_DISPLACE",
|
|
)
|
|
layout.operator(
|
|
"snapform.convert_to_surface",
|
|
text="Convert to Surface",
|
|
icon="MESH_GRID",
|
|
)
|
|
layout.separator()
|
|
layout.operator(
|
|
"snapform.remove_image_background",
|
|
text="Remove Background",
|
|
icon="IMAGE_ALPHA",
|
|
)
|
|
|
|
|
|
def draw_image_context_menu(self, context):
|
|
if not is_image_empty(context.object):
|
|
return
|
|
self.layout.separator()
|
|
self.layout.menu(
|
|
SnapFormImageMenu.bl_idname,
|
|
text="SnapForm",
|
|
icon="PLUGIN",
|
|
)
|