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.*?)(?P\d+)(?P\.[^.]+)$", 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.*?)(?P\d+)(?P\.[^.]+)$", 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", )