- Implement sync_models.py for preparing and uploading ONNX artifacts. - Create build_node.py for building and validating Blender scripts. - Add docs.py for building and previewing documentation with Mike or MkDocs. - Introduce model_exports.py for exporting PyTorch models to ONNX format. - Develop pack.py for packaging the project with dependencies for different platforms. - Add test_sync_models.py to validate sync_models functionality. - Add test_tool_commands.py to ensure repository tools have correct command scripts. - Initialize tools package with __init__.py.
571 lines
18 KiB
Python
571 lines
18 KiB
Python
import re
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import bpy
|
|
|
|
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 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 require_image_empty(object_name):
|
|
"""Return a named Image Empty or raise when it is no longer available."""
|
|
source_object = bpy.data.objects.get(object_name)
|
|
if not is_image_empty(source_object):
|
|
raise RuntimeError("The source Image Empty is no longer available")
|
|
return source_object
|
|
|
|
|
|
def image_pixels(image):
|
|
"""Read one Blender Image into a flat float32 pixel array."""
|
|
import numpy as np
|
|
|
|
pixels = np.empty(len(image.pixels), dtype=np.float32)
|
|
image.pixels.foreach_get(pixels)
|
|
return pixels
|
|
|
|
|
|
def image_rgba(image):
|
|
"""Read one Blender Image into a top-down RGBA float32 array."""
|
|
import numpy as np
|
|
|
|
width, height = (int(value) for value in image.size)
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError("The source image has no readable dimensions")
|
|
return np.flipud(image_pixels(image).reshape((height, width, 4)))
|
|
|
|
|
|
|
|
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 source_start_frame(source_object_name):
|
|
source_object = bpy.data.objects.get(source_object_name)
|
|
return extract_frame_range(source_object)[0] if source_object else 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 cleanup_image_input(path, temporary):
|
|
"""Remove a prepared input and its temporary container when requested."""
|
|
if not temporary or not path:
|
|
return
|
|
path = Path(path)
|
|
try:
|
|
if path.is_dir():
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
else:
|
|
path.unlink(missing_ok=True)
|
|
path.parent.rmdir()
|
|
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.exr"
|
|
|
|
|
|
def normal_data_name(source_object):
|
|
return f"{image_base_name(source_object.data)}_normal.png"
|
|
|
|
|
|
def color_data_name(source_object):
|
|
return f"{image_base_name(source_object.data)}_color.png"
|
|
|
|
|
|
def save_packed_image(image, _scene):
|
|
output_path = Path(tempfile.mkdtemp(prefix="anyimage-input-")) / "input.png"
|
|
original_file_format = image.file_format
|
|
try:
|
|
image.file_format = "PNG"
|
|
image.save(filepath=str(output_path), save_copy=True)
|
|
except (OSError, RuntimeError) as error:
|
|
output_path.unlink(missing_ok=True)
|
|
raise RuntimeError("Unable to export the packed image") from error
|
|
finally:
|
|
image.file_format = original_file_format
|
|
if not output_path.is_file():
|
|
raise RuntimeError("Blender did not export the packed image")
|
|
return output_path
|
|
|
|
|
|
def collect_sequence_frames(image):
|
|
"""Return the numbered frame files of a SEQUENCE image, or [] when absent."""
|
|
source_path = image_path(image)
|
|
if not is_animated_image(image) or not source_path.is_file():
|
|
return []
|
|
match = re.match(
|
|
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
|
|
source_path.name,
|
|
)
|
|
if match is None:
|
|
return []
|
|
prefix = match.group("prefix")
|
|
suffix = match.group("suffix").lower()
|
|
frames = []
|
|
for path in source_path.parent.iterdir():
|
|
frame_match = re.match(
|
|
r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
|
|
path.name,
|
|
)
|
|
if (
|
|
frame_match is not None
|
|
and frame_match.group("prefix") == prefix
|
|
and frame_match.group("suffix").lower() == suffix
|
|
):
|
|
frames.append((int(frame_match.group("digits")), path))
|
|
return [path for _digits, path in sorted(frames)]
|
|
|
|
|
|
def prepare_image_input(image, scene):
|
|
source_path = image_path(image)
|
|
if image_is_packed(image):
|
|
return save_packed_image(image, scene), True
|
|
if getattr(image, "source", None) == "MOVIE":
|
|
if not source_path.is_file():
|
|
raise RuntimeError(
|
|
"The image empty must reference an existing video file"
|
|
)
|
|
return source_path, False
|
|
sequence_frames = collect_sequence_frames(image)
|
|
if sequence_frames:
|
|
temporary_directory = Path(tempfile.mkdtemp(prefix="anyimage-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 prepare_region_input(image, region_mask):
|
|
"""Write one cropped RGBA Region input owned by its temporary directory."""
|
|
import numpy as np
|
|
|
|
left, top, right, bottom = region_mask.bounds
|
|
rgba = image_rgba(image)
|
|
cropped = rgba[top:bottom, left:right].copy()
|
|
if cropped.shape[:2] != region_mask.values.shape:
|
|
raise ValueError("The Region does not match the source image")
|
|
cropped[:, :, 3] *= region_mask.values
|
|
cropped[cropped[:, :, 3] <= 0.0, :3] = 0.0
|
|
output_directory = Path(tempfile.mkdtemp(prefix="anyimage-region-"))
|
|
output_path = output_directory / "input.png"
|
|
temporary_image = bpy.data.images.new(
|
|
"AnyImage Region Input",
|
|
width=right - left,
|
|
height=bottom - top,
|
|
alpha=True,
|
|
)
|
|
original_format = temporary_image.file_format
|
|
try:
|
|
temporary_image.colorspace_settings.name = image.colorspace_settings.name
|
|
temporary_image.alpha_mode = "STRAIGHT"
|
|
temporary_image.pixels.foreach_set(np.flipud(cropped).ravel())
|
|
temporary_image.file_format = "PNG"
|
|
temporary_image.save(filepath=str(output_path), save_copy=True)
|
|
except (OSError, RuntimeError):
|
|
cleanup_image_input(output_path, True)
|
|
raise
|
|
finally:
|
|
temporary_image.file_format = original_format
|
|
bpy.data.images.remove(temporary_image, do_unlink=True)
|
|
if not output_path.is_file():
|
|
cleanup_image_input(output_path, True)
|
|
raise RuntimeError("Blender did not export the Region input")
|
|
return output_path
|
|
|
|
|
|
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 <= 1:
|
|
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 load_normal_result_image(path, source_object):
|
|
"""Load, name, configure, and pack a generated normal image."""
|
|
image = bpy.data.images.load(str(path), check_existing=False)
|
|
image.name = normal_data_name(source_object)
|
|
image.colorspace_settings.name = "Non-Color"
|
|
image.pack()
|
|
return image
|
|
|
|
|
|
def create_pixel_result_image(source_image, pixels, output_size, suffix=""):
|
|
"""Create a packed Image datablock from edited RGBA pixels."""
|
|
output_width, output_height = (int(value) for value in output_size)
|
|
if output_width <= 0 or output_height <= 0:
|
|
raise ValueError("The edited image dimensions must be positive")
|
|
if len(pixels) != output_width * output_height * 4:
|
|
raise ValueError("The edited pixels do not match the image dimensions")
|
|
image = bpy.data.images.new(
|
|
f"{image_base_name(source_image)}{suffix}",
|
|
width=output_width,
|
|
height=output_height,
|
|
alpha=True,
|
|
)
|
|
try:
|
|
image.colorspace_settings.name = source_image.colorspace_settings.name
|
|
image.alpha_mode = "STRAIGHT"
|
|
image.pixels.foreach_set(pixels)
|
|
image.pack()
|
|
image.update()
|
|
except Exception:
|
|
bpy.data.images.remove(image, do_unlink=True)
|
|
raise
|
|
return image
|
|
|
|
|
|
def replace_empty_image(source_object, result_image, *, placement_bounds=None):
|
|
"""Replace one Empty image, retaining shared source images."""
|
|
source_image = source_object.data
|
|
source_name = source_image.name
|
|
original_fake_user = source_image.use_fake_user
|
|
image_user = getattr(source_object, "image_user", None)
|
|
frame_count = (
|
|
video_total_frames(result_image)
|
|
if is_animated_image(result_image)
|
|
else None
|
|
)
|
|
alignment = (
|
|
empty_display_alignment(
|
|
source_object,
|
|
tuple(source_image.size),
|
|
tuple(result_image.size),
|
|
placement_bounds,
|
|
)
|
|
if placement_bounds is not None
|
|
else None
|
|
)
|
|
if frame_count is not None:
|
|
(
|
|
original_frame_start,
|
|
_original_frame_offset,
|
|
_original_frame_duration,
|
|
original_use_cyclic,
|
|
) = image_user_settings(source_object)
|
|
|
|
source_image.use_fake_user = True
|
|
try:
|
|
source_object.data = result_image
|
|
if frame_count is not None:
|
|
configure_image_user(
|
|
image_user,
|
|
frame_count,
|
|
frame_start=original_frame_start,
|
|
use_cyclic=original_use_cyclic,
|
|
)
|
|
if alignment is not None:
|
|
source_object.empty_display_size = alignment[0]
|
|
source_object.empty_image_offset = alignment[1]
|
|
except Exception:
|
|
source_object.data = source_image
|
|
source_image.use_fake_user = original_fake_user
|
|
raise
|
|
|
|
source_image.use_fake_user = original_fake_user
|
|
if not original_fake_user and int(source_image.users) == 0:
|
|
bpy.data.images.remove(source_image, do_unlink=True)
|
|
result_image.name = source_name
|
|
return result_image
|
|
|
|
|
|
|
|
def load_color_result_image(
|
|
source_object,
|
|
output_directory,
|
|
*,
|
|
frame_prefix="image.",
|
|
):
|
|
output_directory = Path(output_directory)
|
|
color_paths = sorted_frame_paths(output_directory, frame_prefix)
|
|
if not color_paths:
|
|
raise RuntimeError(
|
|
f"Color result has no images: {output_directory}"
|
|
)
|
|
source_image = source_object.data
|
|
source_color_space = source_image.colorspace_settings.name
|
|
animated = is_animated_image(source_image)
|
|
try:
|
|
result_image = bpy.data.images.load(
|
|
str(color_paths[0]),
|
|
check_existing=False,
|
|
)
|
|
result_image.alpha_mode = "CHANNEL_PACKED"
|
|
result_image.colorspace_settings.name = source_color_space
|
|
if animated:
|
|
force_sequence_source(result_image)
|
|
except (OSError, RuntimeError, TypeError, ValueError) 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
|
|
|
|
try:
|
|
result_name = color_data_name(source_object)
|
|
if animated:
|
|
pack_sequence_if_possible(result_image, len(color_paths))
|
|
else:
|
|
result_image.pack()
|
|
result_image.name = result_name
|
|
except (OSError, RuntimeError, TypeError, ValueError) as error:
|
|
bpy.data.images.remove(result_image, do_unlink=True)
|
|
raise RuntimeError("Unable to prepare the color result") from error
|
|
return result_image
|
|
|
|
|
|
def image_empty_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 empty_display_alignment(
|
|
source_object,
|
|
source_size,
|
|
output_size,
|
|
placement_bounds,
|
|
):
|
|
"""Return the display size and offset for an edited Image Empty."""
|
|
source_width, source_height = source_size
|
|
output_width, output_height = output_size
|
|
left, top, right, bottom = placement_bounds
|
|
if (
|
|
source_width <= 0
|
|
or source_height <= 0
|
|
or output_width <= 0
|
|
or output_height <= 0
|
|
or right <= left
|
|
or bottom <= top
|
|
):
|
|
raise ValueError("The edited image placement is invalid")
|
|
x_min, x_max, y_min, y_max = image_empty_bounds(source_object)
|
|
local_width = x_max - x_min
|
|
local_height = y_max - y_min
|
|
desired_x_min = x_min + left * local_width / source_width
|
|
desired_x_max = x_min + right * local_width / source_width
|
|
desired_y_max = y_max - top * local_height / source_height
|
|
desired_y_min = y_max - bottom * local_height / source_height
|
|
desired_width = desired_x_max - desired_x_min
|
|
desired_height = desired_y_max - desired_y_min
|
|
display_size = max(desired_width, desired_height)
|
|
aspect = output_width / output_height
|
|
if aspect >= 1.0:
|
|
display_width = display_size
|
|
display_height = display_size / aspect
|
|
else:
|
|
display_width = display_size * aspect
|
|
display_height = display_size
|
|
center_x = (desired_x_min + desired_x_max) * 0.5
|
|
center_y = (desired_y_min + desired_y_max) * 0.5
|
|
return display_size, (
|
|
center_x / display_width - 0.5,
|
|
center_y / display_height - 0.5,
|
|
)
|