Refactor image handling and update foreground processing

- Changed the logic in `remove_background_faces` to keep only faces with all foreground vertices.
- Updated background job to save foreground images with a new naming convention (`color.0001.png`).
- Replaced `copy_frames` with `copy_color_images` in various jobs to ensure consistent image naming.
- Introduced a new function `copy_color_images` to handle the copying of images with a specific naming format.
- Updated the BEN2 model to save output images with a new prefix for better clarity.
- Modified tests to reflect changes in image naming and processing logic.
- Added a new shadeless node group for improved material handling in geometry nodes.
- Enhanced validation scripts to check for new image naming conventions and ensure proper functionality.
This commit is contained in:
2026-08-19 01:45:59 +08:00
parent 835225d264
commit 20b653bb30
28 changed files with 663 additions and 383 deletions

Binary file not shown.

View File

@@ -23,6 +23,8 @@ 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"
SHADELESS_NODE_GROUP_NAME = "O Shadeless"
MATERIAL_COLOR_SPACE = "AgX Base sRGB"
DA3_PATCH_SIZE = 14
DA3_MIN_PROCESS_RES = 252
DA3_MAX_PROCESS_RES = 1008
@@ -183,8 +185,8 @@ def copy_depth_image(source_object):
return depth_image
def removed_background_data_name(source_object):
return f"{image_base_name(source_object.data)}.foreground"
def color_data_name(source_object):
return f"{image_base_name(source_object.data)}.color"
def scaled_process_resolution(image, scale_percent):
@@ -328,7 +330,7 @@ def configure_image_user(
def pack_sequence_if_possible(image, expected_frames):
"""Pack a frame sequence when Blender supports it; otherwise keep external files."""
if expected_frames <= 0:
if expected_frames <= 1:
return False
try:
image.pack()
@@ -455,81 +457,60 @@ def create_depth_sequence_image(name, depth_frames, scene, frame_directory):
return sequence_image
def create_removed_background_image(
def create_color_image(
source_object,
output_path,
output_directory,
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():
output_directory = Path(output_directory)
color_paths = sorted_frame_paths(output_directory, "color.")
if not color_paths:
raise RuntimeError(
f"Background removal result does not exist: {output_path}"
f"Color result has no images: {output_directory}"
)
source_image = source_object.data
animated = is_animated_image(source_image)
try:
result_image = bpy.data.images.load(
str(output_path),
str(color_paths[0]),
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
if animated:
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 color result") from error
original_fake_user = source_image.use_fake_user
try:
result_image.name = removed_background_data_name(source_object)
result_image.pack()
result_image.name = color_data_name(source_object)
if animated:
pack_sequence_if_possible(result_image, len(color_paths))
else:
result_image.pack()
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
if animated:
configure_image_user(
image_user,
len(color_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 pack the background result") from error
bpy.data.images.remove(result_image, do_unlink=True)
raise RuntimeError("Unable to prepare the color result") from error
return result_image
@@ -993,7 +974,7 @@ def create_moge_material(
color_image,
depth_image,
depth_scale,
unshaded=False,
shadeless=False,
frame_count=1,
frame_start=1,
use_cyclic=False,
@@ -1009,6 +990,7 @@ def create_moge_material(
uv.uv_map = "UVMap"
uv.location = (-900.0, 200.0)
color = nodes.new("ShaderNodeTexImage")
color_image.colorspace_settings.name = MATERIAL_COLOR_SPACE
color.image = color_image
color.interpolation = "Linear"
color.extension = "EXTEND"
@@ -1043,14 +1025,16 @@ def create_moge_material(
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")
if shadeless:
shader = nodes.new("ShaderNodeGroup")
shader.node_tree = load_shadeless_node_group()
shader.name = SHADELESS_NODE_GROUP_NAME
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(shader.outputs["Shader"], mix.inputs[2])
links.new(mix.outputs["Shader"], output.inputs["Surface"])
return material
shader = nodes.new("ShaderNodeBsdfPrincipled")
@@ -1307,16 +1291,13 @@ def create_material(
color_image,
depth_scale,
displacement_image=None,
foreground_image=None,
displacement_midlevel=0.0,
displacement_scale=1.0,
unshaded=False,
shadeless=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,
@@ -1336,7 +1317,8 @@ def create_material(
uv_map.uv_map = "UVMap"
uv_map.location = (-1250.0, 350.0)
material_image = foreground_image or color_image
material_image = color_image
material_image.colorspace_settings.name = MATERIAL_COLOR_SPACE
color_texture.image = material_image
color_texture.interpolation = "Linear"
color_texture.extension = "EXTEND"
@@ -1355,13 +1337,6 @@ def create_material(
)
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,
@@ -1394,16 +1369,18 @@ def create_material(
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")
if shadeless:
shader = nodes.new("ShaderNodeGroup")
shader.node_tree = load_shadeless_node_group()
shader.name = SHADELESS_NODE_GROUP_NAME
transparent = nodes.new("ShaderNodeBsdfTransparent")
mix_shader = nodes.new("ShaderNodeMixShader")
emission.location = (0.0, 550.0)
shader.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"],
shader.inputs["Color"],
)
links.new(
material_inputs.outputs["Alpha"],
@@ -1414,7 +1391,7 @@ def create_material(
mix_shader.inputs[1],
)
links.new(
emission.outputs["Emission"],
shader.outputs["Shader"],
mix_shader.inputs[2],
)
links.new(
@@ -1502,13 +1479,13 @@ 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):
def load_asset_node_group(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}")
raise RuntimeError(f"Missing node asset: {asset_path}")
use_pack = bpy.app.version >= (5, 0, 0)
if not use_pack:
directory = str(asset_path) + "/NodeTree/"
@@ -1520,18 +1497,15 @@ def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME):
do_reuse_local_id=True,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to append the Geometry Nodes asset")
raise RuntimeError("Unable to append the node asset")
node_group = bpy.data.node_groups.get(group_name)
if node_group is None:
raise RuntimeError("Appended Geometry Nodes asset is unavailable")
raise RuntimeError("Appended node 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}"
)
raise RuntimeError(f"Node asset is missing {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)
@@ -1541,6 +1515,14 @@ def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME):
return node_group
def load_shadeless_node_group():
return load_asset_node_group(SHADELESS_NODE_GROUP_NAME)
def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME):
return load_asset_node_group(group_name)
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":
@@ -1600,7 +1582,7 @@ def drive_modifier_input(owner_object, modifier, identifier, expression="frame")
def add_depth_plane_processing(
plane_object,
depth_image,
foreground_image,
color_image,
material,
frame_start=1,
frame_duration=1,
@@ -1616,7 +1598,7 @@ def add_depth_plane_processing(
modifier.show_group_selector = False
for name, value in (
("Depth Image", depth_image),
("Foreground Image", foreground_image),
("Color Image", color_image),
("Material", material),
("Frame", 1),
("Frame Start", frame_start),
@@ -1635,6 +1617,7 @@ def add_depth_plane_processing(
def add_moge_depth_plane_processing(
plane_object,
depth_image,
color_image,
material,
frame_start=1,
frame_duration=1,
@@ -1652,6 +1635,7 @@ def add_moge_depth_plane_processing(
modifier.show_group_selector = False
for name, value in (
("Depth Image", depth_image),
("Color Image", color_image),
("Material", material),
("Frame", 1),
("Frame Start", frame_start),
@@ -1688,9 +1672,8 @@ def build_image_depth_plane(
context,
source_object,
delete_source,
unshaded=False,
shadeless=False,
color_image=None,
foreground_image=None,
depth_frames_directory=None,
depth_image=None,
):
@@ -1701,9 +1684,10 @@ def build_image_depth_plane(
frames_image = create_sequence_image(
depth_data_name(source_object),
depth_frames_directory,
frame_prefix="color.",
)
frame_count = len(
sorted_frame_paths(depth_frames_directory, "frame.")
sorted_frame_paths(depth_frames_directory, "color.")
)
(
source_frame_start,
@@ -1722,9 +1706,9 @@ def build_image_depth_plane(
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
)
material_image = color_image
if material_image is None:
raise RuntimeError("Depth Plane material image is required")
mesh = create_plane_mesh(source_object, subdivision_level=None)
plane_object = bpy.data.objects.new(source_name, mesh)
context.collection.objects.link(plane_object)
@@ -1733,19 +1717,16 @@ def build_image_depth_plane(
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,
material_image,
depth_scale,
displacement_image=depth_source_image,
foreground_image=foreground_image,
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
unshaded=unshaded,
shadeless=shadeless,
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,
@@ -1754,7 +1735,7 @@ def build_image_depth_plane(
add_depth_plane_processing(
plane_object,
depth_source_image,
foreground_image,
material_image,
material,
frame_start=source_frame_start,
frame_duration=frame_count,
@@ -1775,10 +1756,9 @@ def build_depth_plane(
source_object,
npz_path,
delete_source,
unshaded=False,
shadeless=False,
depth_name="",
color_image=None,
foreground_image=None,
depth_frames_directory=None,
):
source_name = source_object.name
@@ -1797,9 +1777,9 @@ def build_depth_plane(
]
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
)
material_image = color_image
if material_image is None:
raise RuntimeError("Depth Plane material image is required")
(
source_frame_start,
source_frame_offset,
@@ -1825,7 +1805,10 @@ def build_depth_plane(
depth_name,
scaled_depth_arrays[0],
)
if frame_count > 1 and depth_frames_directory is not None:
if (
is_animated_image(source_object.data)
and depth_frames_directory is not None
):
bpy.data.images.remove(depth_image, do_unlink=True)
depth_image = create_depth_sequence_image(
depth_name,
@@ -1836,21 +1819,18 @@ def build_depth_plane(
material = create_material(
source_name,
color_image if color_image is not None else source_object.data,
material_image,
depth_scale,
displacement_image=depth_image,
foreground_image=foreground_image,
displacement_midlevel=DISPLACEMENT_MIDLEVEL,
displacement_scale=DISPLACEMENT_STRENGTH,
unshaded=unshaded,
shadeless=shadeless,
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,
@@ -1859,7 +1839,7 @@ def build_depth_plane(
add_depth_plane_processing(
plane_object,
depth_image,
foreground_image,
material_image,
material,
frame_start=source_frame_start,
frame_duration=frame_count,
@@ -1885,9 +1865,8 @@ def build_moge2_depth_plane(
source_object,
npz_path,
delete_source,
unshaded=False,
shadeless=False,
color_image=None,
foreground_image=None,
data_frames_directory=None,
):
import numpy as np
@@ -1935,7 +1914,10 @@ def build_moge2_depth_plane(
depth_data_name(source_object),
displacement[0] * depth_scale,
)
if len(depth) > 1 and data_frames_directory is not None:
if (
is_animated_image(source_object.data)
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),
@@ -1945,12 +1927,15 @@ def build_moge2_depth_plane(
"moge_depth",
"depth",
)
material_image = color_image
if material_image is None:
raise RuntimeError("Depth Plane material image is required")
material = create_moge_material(
source_object.name,
foreground_image or color_image or source_object.data,
material_image,
depth_image,
depth_scale,
unshaded=unshaded,
shadeless=shadeless,
frame_count=len(depth),
frame_start=source_frame_start,
use_cyclic=source_use_cyclic,
@@ -1959,6 +1944,7 @@ def build_moge2_depth_plane(
modifier = add_moge_depth_plane_processing(
plane_object,
depth_image,
material_image,
material,
frame_start=source_frame_start,
frame_duration=len(depth),
@@ -2028,9 +2014,9 @@ class ConvertToDepthPlane(bpy.types.Operator):
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",
shadeless: bpy.props.BoolProperty(
name="Shadeless",
description="Display the source color without scene lighting",
default=False,
)
remove_background: bpy.props.BoolProperty(
@@ -2081,7 +2067,7 @@ class ConvertToDepthPlane(bpy.types.Operator):
elif self.model.startswith("MOGE2_"):
layout.prop(self, "resolution_level")
layout.prop(self, "remove_background")
layout.prop(self, "unshaded")
layout.prop(self, "shadeless")
layout.separator(type="LINE")
layout.prop(self, "keep_original")
@@ -2110,7 +2096,7 @@ class ConvertToDepthPlane(bpy.types.Operator):
source_object_name=source_object.name,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
shadeless=self.shadeless,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2134,7 +2120,7 @@ class ConvertToDepthPlane(bpy.types.Operator):
model=self.model,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
shadeless=self.shadeless,
remove_background=self.remove_background,
)
if "CANCELLED" not in result:
@@ -2169,7 +2155,7 @@ class ConvertToDepthPlane(bpy.types.Operator):
process_res=process_res,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
shadeless=self.shadeless,
remove_background=self.remove_background,
)
if "CANCELLED" in result:

View File

@@ -402,7 +402,7 @@ def build_image_mesh(
npz_path,
delete_source,
create_camera,
unshaded,
shadeless,
moge2=False,
color_image=None,
normal_path=None,
@@ -437,6 +437,8 @@ def build_image_mesh(
context.collection.objects.link(mesh_object)
mesh_object.matrix_world = source_object.matrix_world.copy()
depth_scale = material_depth_scale(mesh_object, camera_parameters)
if color_image is None:
raise RuntimeError("Surface material image is required")
normal_image = (
load_surface_normal_image(
normal_path,
@@ -447,9 +449,9 @@ def build_image_mesh(
)
material = create_material(
source_name,
color_image or source_object.data,
color_image,
depth_scale,
unshaded=unshaded,
shadeless=shadeless,
surface_group_name=(
MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME if moge2 else None
),
@@ -509,9 +511,9 @@ class ConvertToSurface(bpy.types.Operator):
description="Create a perspective camera matching the source projection",
default=False,
)
unshaded: bpy.props.BoolProperty(
name="Unshaded",
description="Display the source color through an Emission shader",
shadeless: bpy.props.BoolProperty(
name="Shadeless",
description="Display the source color without scene lighting",
default=False,
)
remove_background: bpy.props.BoolProperty(
@@ -550,7 +552,7 @@ class ConvertToSurface(bpy.types.Operator):
draw_process_resolution(layout, self, context.object.data)
layout.prop(self, "create_camera")
layout.prop(self, "remove_background")
layout.prop(self, "unshaded")
layout.prop(self, "shadeless")
layout.separator(type="LINE")
layout.prop(self, "keep_original")
@@ -580,7 +582,7 @@ class ConvertToSurface(bpy.types.Operator):
"delete_source": not self.keep_original,
"delete_input": delete_input,
"create_camera": self.create_camera,
"unshaded": self.unshaded,
"shadeless": self.shadeless,
"remove_background": self.remove_background,
}
if self.model.startswith("MOGE2_"):

View File

@@ -111,13 +111,12 @@ def build_da3_depth_response(
depth_data_name,
remove_background,
delete_source,
unshaded,
shadeless,
):
from .blender.depth_plane import (
build_depth_plane,
create_removed_background_image,
create_sequence_image,
image_base_name,
create_color_image,
is_animated_image,
is_image_empty,
)
@@ -125,32 +124,21 @@ def build_da3_depth_response(
source_object = bpy.data.objects.get(source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source image empty is no longer available")
foreground_paths = sorted(directory.glob("foreground.*.png"))
foreground_image = None
if remove_background:
foreground_image = create_removed_background_image(
source_object,
directory if foreground_paths else directory / "foreground.png",
assign_to_source=False,
)
color_image = None
if list(directory.glob("frame.*")) and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
directory,
)
animated = is_animated_image(source_object.data)
color_image = create_color_image(
source_object,
directory,
assign_to_source=False,
)
build_depth_plane(
context,
source_object,
result.file("depth"),
delete_source,
unshaded=unshaded,
shadeless=shadeless,
depth_name=depth_data_name,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
directory if list(directory.glob("frame.*")) else None
),
depth_frames_directory=directory if animated else None,
)
push_undo("Convert to Depth Plane")
@@ -162,13 +150,11 @@ def build_moge2_depth_response(
source_object_name,
remove_background,
delete_source,
unshaded,
shadeless,
):
from .blender.depth_plane import (
build_moge2_depth_plane,
create_removed_background_image,
create_sequence_image,
image_base_name,
create_color_image,
is_image_empty,
)
@@ -176,28 +162,18 @@ def build_moge2_depth_response(
source_object = bpy.data.objects.get(source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source image empty is no longer available")
foreground_paths = sorted(directory.glob("foreground.*.png"))
foreground_image = None
if remove_background:
foreground_image = create_removed_background_image(
source_object,
directory if foreground_paths else directory / "foreground.png",
assign_to_source=False,
)
color_image = None
if list(directory.glob("frame.*")) and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
directory,
)
color_image = create_color_image(
source_object,
directory,
assign_to_source=False,
)
build_moge2_depth_plane(
context,
source_object,
result.file("depth"),
delete_source,
unshaded=unshaded,
shadeless=shadeless,
color_image=color_image,
foreground_image=foreground_image,
data_frames_directory=directory,
)
push_undo("Convert to Depth Plane")
@@ -210,13 +186,12 @@ def build_original_depth_response(
source_object_name,
background,
delete_source,
unshaded,
shadeless,
):
from .blender.depth_plane import (
build_image_depth_plane,
create_removed_background_image,
create_sequence_image,
image_base_name,
create_color_image,
is_animated_image,
is_image_empty,
)
@@ -224,30 +199,19 @@ def build_original_depth_response(
source_object = bpy.data.objects.get(source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source image empty is no longer available")
foreground_image = None
if background:
foreground_paths = sorted(directory.glob("foreground.*.png"))
foreground_image = create_removed_background_image(
source_object,
directory if foreground_paths else directory / "foreground.png",
assign_to_source=False,
)
color_image = None
if list(directory.glob("frame.*")) and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
directory,
)
animated = is_animated_image(source_object.data)
color_image = create_color_image(
source_object,
directory,
assign_to_source=False,
)
build_image_depth_plane(
context,
source_object,
delete_source,
unshaded=unshaded,
shadeless=shadeless,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
directory if list(directory.glob("frame.*")) else None
),
depth_frames_directory=directory if animated else None,
)
push_undo("Convert to Depth Plane")
@@ -259,27 +223,19 @@ def build_surface_response(
source_object_name,
delete_source,
create_camera,
unshaded,
shadeless,
moge2,
):
from .blender.depth_plane import (
create_removed_background_image,
is_image_empty,
)
from .blender.depth_plane import create_color_image, is_image_empty
from .blender.surface import build_image_mesh
source_object = bpy.data.objects.get(source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source Image Empty no longer exists")
foreground_path = result.directory / "foreground.png"
color_image = (
create_removed_background_image(
source_object,
foreground_path,
assign_to_source=False,
)
if foreground_path.is_file()
else None
color_image = create_color_image(
source_object,
result.directory,
assign_to_source=False,
)
build_image_mesh(
context,
@@ -287,7 +243,7 @@ def build_surface_response(
result.file("surface"),
delete_source,
create_camera,
unshaded,
shadeless,
moge2,
color_image=color_image,
normal_path=(
@@ -300,19 +256,12 @@ def build_surface_response(
def apply_background_response(result, source_object_name):
from .blender.depth_plane import (
create_removed_background_image,
is_image_empty,
)
from .blender.depth_plane import create_color_image, is_image_empty
source_object = bpy.data.objects.get(source_object_name)
if not is_image_empty(source_object):
raise RuntimeError("The source image empty is no longer available")
foreground_paths = sorted(result.directory.glob("foreground.*.png"))
create_removed_background_image(
source_object,
result.directory if foreground_paths else result.file("foreground"),
)
create_color_image(source_object, result.directory)
push_undo("Remove Background")
@@ -461,7 +410,7 @@ class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator):
process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN)
delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN)
shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN)
remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN)
def request(self, context):
@@ -492,7 +441,7 @@ class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator):
depth_data_name=self.depth_data_name,
remove_background=self.remove_background,
delete_source=self.delete_source,
unshaded=self.unshaded,
shadeless=self.shadeless,
)
return "Depth plane is ready"
@@ -512,7 +461,7 @@ class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator):
resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN)
delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN)
shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN)
remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN)
def request(self, context):
@@ -541,7 +490,7 @@ class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator):
source_object_name=self.source_object_name,
remove_background=self.remove_background,
delete_source=self.delete_source,
unshaded=self.unshaded,
shadeless=self.shadeless,
)
return "MoGe-2 depth plane is ready"
@@ -562,7 +511,7 @@ class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator):
delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN)
unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN)
shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN)
remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN)
def request(self, context):
@@ -591,7 +540,7 @@ class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator):
source_object_name=self.source_object_name,
delete_source=self.delete_source,
create_camera=self.create_camera,
unshaded=self.unshaded,
shadeless=self.shadeless,
moge2=False,
)
return "Perspective surface is ready"
@@ -613,7 +562,7 @@ class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator):
delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN)
unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN)
shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN)
remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN)
def request(self, context):
@@ -642,7 +591,7 @@ class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator):
source_object_name=self.source_object_name,
delete_source=self.delete_source,
create_camera=self.create_camera,
unshaded=self.unshaded,
shadeless=self.shadeless,
moge2=True,
)
return "MoGe-2 surface is ready"
@@ -691,7 +640,7 @@ class GenerateOriginalBackground(JobOperatorBase, bpy.types.Operator):
source_object_name: bpy.props.StringProperty(options=HIDDEN)
delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN)
shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN)
def request(self, context):
require_environment()
@@ -712,7 +661,7 @@ class GenerateOriginalBackground(JobOperatorBase, bpy.types.Operator):
source_object_name=self.source_object_name,
background=True,
delete_source=self.delete_source,
unshaded=self.unshaded,
shadeless=self.shadeless,
)
return "Depth plane is ready"
@@ -730,7 +679,7 @@ class ExtractOriginalFrames(JobOperatorBase, bpy.types.Operator):
source_object_name: bpy.props.StringProperty(options=HIDDEN)
delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN)
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN)
shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN)
def request(self, context):
require_environment()
@@ -748,7 +697,7 @@ class ExtractOriginalFrames(JobOperatorBase, bpy.types.Operator):
source_object_name=self.source_object_name,
background=False,
delete_source=self.delete_source,
unshaded=self.unshaded,
shadeless=self.shadeless,
)
return "Depth plane is ready"

View File

@@ -69,9 +69,9 @@ def triangle_edge_stretch_values(points, triangles):
def remove_background_faces(faces, foreground_mask):
"""Keep faces with at least one foreground vertex."""
"""Keep only faces whose vertices are all foreground."""
foreground_points = foreground_mask.reshape(-1)
return faces[foreground_points[faces].any(axis=1)]
return faces[foreground_points[faces].all(axis=1)]
def depth_edge_mask(depth, threshold=12.0 / 255.0):

View File

@@ -1,7 +1,7 @@
import shutil
import time
from ..media import copy_frames, prepare_background_input
from ..media import prepare_background_input
from ..models.ben2 import infer_foreground, infer_foreground_frames
@@ -39,7 +39,7 @@ def run(context, parameters):
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
output_path = output_directory / "color.0001.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
@@ -58,8 +58,8 @@ def run(context, parameters):
0.7,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
)
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -73,6 +73,4 @@ def run(context, parameters):
)
print(message)
context.progress(1.0, message)
if tuple(output_directory.glob("foreground.*.png")):
return {"foreground": "."}
return {"foreground": "foreground.png"}
return {"color": "."}

View File

@@ -1,4 +1,4 @@
from ..media import copy_frames
from ..media import copy_color_images
from ..models import da3, geometry_adapter, moge2
from ..outputs.depth import write_da3_depth, write_moge_depth
from .common import input_frames
@@ -36,14 +36,15 @@ def run(context, parameters):
cancel_check=context.check_cancelled,
)
write_moge_depth(context.directory, prediction)
write_optional(
color_paths = write_optional(
context,
parameters,
images,
progress_start=0.84,
progress_span=0.08,
)
copy_frames(images, context.directory)
if not color_paths:
copy_color_images(images, context.directory)
context.check_cancelled()
context.progress(1.0, f"Exported depth for {len(images)} frame(s)")
return {"depth": "results.npz"}

View File

@@ -24,7 +24,7 @@ def write_optional(
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
path = context.directory / "foreground.png"
path = context.directory / "color.0001.png"
foreground.convert("RGBA").save(
path,
format="PNG",
@@ -42,8 +42,9 @@ def write_optional(
progress_span,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
)
paths = sorted(context.directory.glob("foreground.*.png"))
paths = sorted(context.directory.glob("color.*.png"))
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"

View File

@@ -1,6 +1,6 @@
import shutil
from ..media import copy_frames, prepare_input
from ..media import copy_color_images, prepare_input
def run(context, parameters):
@@ -13,7 +13,7 @@ def run(context, parameters):
cancel_check=context.check_cancelled,
)
try:
copy_frames(images, context.directory)
copy_color_images(images, context.directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

View File

@@ -1,15 +1,18 @@
from ..geometry import foreground_mesh_mask
from ..geometry import build_surface, foreground_mesh_mask
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from ..geometry import build_surface
from ..outputs.surface import (
write_normal_texture,
write_surface_archive,
write_surface_texture,
)
from ..outputs.surface import write_normal_texture, write_surface_archive
from .common import input_frames
from .foreground import write_optional
def inference_progress_message(model_name, index, total):
message = f"Estimating {model_name} geometry"
if total > 1:
message += f" ({index + 1}/{total})"
return message
def run(context, parameters):
context.progress(0.02, "Preparing surface input")
context.directory.mkdir(parents=True, exist_ok=True)
@@ -27,7 +30,7 @@ def run(context, parameters):
include_points=True,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
f"Estimating MoGe-2 geometry ({index + 1}/{total})",
inference_progress_message("MoGe-2", index, total),
),
cancel_check=context.check_cancelled,
)
@@ -38,13 +41,13 @@ def run(context, parameters):
images,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
f"Estimating DA3 geometry ({index + 1}/{total})",
inference_progress_message("DA3", index, total),
),
cancel_check=context.check_cancelled,
)
context.check_cancelled()
frame = adapter.surface_frame(prediction)
foreground_paths = write_optional(
color_paths = write_optional(
context,
parameters,
images,
@@ -52,19 +55,17 @@ def run(context, parameters):
progress_span=0.08,
)
foreground_mask = None
if foreground_paths:
if color_paths:
foreground_mask = foreground_mesh_mask(
foreground_paths[0],
color_paths[0],
frame.depth.shape,
)
context.progress(0.82, "Building surface")
mesh = build_surface(frame, foreground_mask=foreground_mask)
context.progress(0.92, "Writing surface")
write_surface_archive(mesh, context.directory / "mesh.npz")
write_surface_texture(
foreground_paths[0] if foreground_paths else images[0],
context.directory / "texture.png",
)
if not color_paths:
copy_color_images(images, context.directory)
if adapter is moge2:
write_normal_texture(
prediction.normal[0],

View File

@@ -4,6 +4,7 @@ from .input import (
TaskCancelled,
check_cancelled,
collect_images,
copy_color_images,
copy_frames,
extract_video_frames,
first_frames,
@@ -18,6 +19,7 @@ __all__ = (
"TaskCancelled",
"check_cancelled",
"collect_images",
"copy_color_images",
"copy_frames",
"extract_video_frames",
"first_frames",

View File

@@ -143,16 +143,26 @@ def prepare_background_input(
)
def copy_frames(images, output_directory):
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
if len(images) <= 1:
return
def copy_color_images(images, output_directory):
"""Copy input images as the canonical numbered color layer."""
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
Path(output_directory) / f"color.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)
def copy_frames(images, output_directory):
"""Copy multiple input images as numbered raw Debug frame files."""
if len(images) <= 1:
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)

View File

@@ -48,8 +48,9 @@ def infer_foreground_frames(
progress_span,
model_cache=None,
cancel_check=None,
output_prefix="foreground",
):
"""Run BEN2 on every frame and write foreground.0001.png etc."""
"""Run BEN2 on every frame and write numbered transparent PNG files."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
@@ -65,7 +66,7 @@ def infer_foreground_frames(
model_cache=model_cache,
cancel_check=cancel_check,
)
output_path = output_directory / f"foreground.{index + 1:04d}.png"
output_path = output_directory / f"{output_prefix}.{index + 1:04d}.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",

View File

@@ -1,5 +1,5 @@
from .depth import write_da3_depth, write_moge_depth
from .surface import write_normal_texture, write_surface_archive, write_surface_texture
from .surface import write_normal_texture, write_surface_archive
__all__ = (
@@ -7,5 +7,4 @@ __all__ = (
"write_moge_depth",
"write_normal_texture",
"write_surface_archive",
"write_surface_texture",
)

View File

@@ -20,15 +20,6 @@ def write_surface_archive(mesh, output_path):
return output_path
def write_surface_texture(image_path, output_path):
from PIL import Image
output_path = Path(output_path)
with Image.open(image_path) as image:
image.convert("RGBA").save(output_path, format="PNG", compress_level=1)
return output_path
def write_normal_texture(normal, output_path):
import numpy as np
from PIL import Image