Add new node asset scripts for geometry and shader operations

- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters.
- Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections.
- Added `validate.py` to validate the integrity of the newly created node groups and their inputs.
- Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies.
- Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
This commit is contained in:
2026-08-30 00:17:29 +08:00
parent 7d39fe8c5e
commit 32bcd45d7b
168 changed files with 7213 additions and 13121 deletions

View File

@@ -1,37 +0,0 @@
import shutil
from contextlib import contextmanager
from ..media import prepare_input
def frame_progress_message(message, index, total):
if total > 1:
return f"{message} ({index + 1}/{total})"
return message
@contextmanager
def input_frames(context, parameters):
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
context.directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
yield images
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
@contextmanager
def analysis_frames(context, parameters, fallback_images):
analysis_input = parameters.get("analysis_input")
if not analysis_input or analysis_input == parameters["input"]:
yield fallback_images
return
analysis_parameters = dict(parameters, input=analysis_input)
with input_frames(context, analysis_parameters) as images:
yield images

View File

@@ -1,176 +0,0 @@
from pathlib import Path
PLANE_VARIANCE_FLOOR_RATIO = 0.025
def resized_selection_mask(selection_mask, shape):
import numpy as np
from PIL import Image
selection_mask = np.asarray(selection_mask, dtype=bool)
if selection_mask.shape == tuple(shape):
return selection_mask
resized = Image.fromarray(selection_mask.astype(np.uint8) * 255, mode="L")
resized = resized.resize(
(int(shape[1]), int(shape[0])),
Image.Resampling.NEAREST,
)
return np.asarray(resized, dtype=np.uint8) > 0
def _stable_weighted_plane(coordinates, values, weights):
import numpy as np
total = float(np.sum(weights))
if total <= 1e-12:
return np.zeros(3, dtype=np.float64)
center = np.sum(coordinates * weights[:, None], axis=0) / total
value_center = float(np.sum(values * weights) / total)
centered = coordinates - center
covariance = (centered * weights[:, None]).T @ centered / total
cross = centered.T @ (weights * (values - value_center)) / total
eigenvalues, eigenvectors = np.linalg.eigh(covariance)
largest = max(float(eigenvalues[-1]), 0.0)
if largest <= 1e-12:
slopes = np.zeros(2, dtype=np.float64)
else:
supported = np.maximum(eigenvalues, largest * PLANE_VARIANCE_FLOOR_RATIO)
slopes = eigenvectors @ ((eigenvectors.T @ cross) / supported)
intercept = value_center - float(center @ slopes)
return np.asarray((*slopes, intercept), dtype=np.float64)
def normalized_moge_depth(depth, model_mask, selection_mask):
"""Normalize MoGe depth while preserving its adjustable fitted base plane."""
import numpy as np
depth = np.asarray(depth, dtype=np.float32)
model_mask = np.asarray(model_mask, dtype=bool)
selection_mask = resized_selection_mask(selection_mask, depth.shape)
if model_mask.shape != depth.shape:
raise ValueError("MoGe-2 depth and mask dimensions do not match")
valid = selection_mask & model_mask & np.isfinite(depth) & (depth > 0.0)
if not valid.any():
raise ValueError("MoGe-2 found no valid depth inside the lasso")
height, width = depth.shape
y, x = np.indices(depth.shape, dtype=np.float64)
x = x / max(width - 1, 1) * 2.0 - 1.0
y = y / max(height - 1, 1) * 2.0 - 1.0
coordinates = np.stack((x[valid], y[valid]), axis=1)
selected = depth[valid].astype(np.float64)
weights = np.ones_like(selected)
coefficients = np.zeros(3, dtype=np.float64)
for _iteration in range(4):
coefficients = _stable_weighted_plane(
coordinates,
selected,
weights,
)
residual = (
selected
- coordinates @ coefficients[:2]
- coefficients[2]
)
centered = residual - np.median(residual)
scale = max(float(np.median(np.abs(centered))) * 1.4826, 1e-6)
weights = np.minimum(
1.0,
(2.5 * scale) / np.maximum(np.abs(centered), 1e-12),
)
baseline = coefficients[0] * x + coefficients[1] * y + coefficients[2]
reference = float(np.median(baseline[valid]))
residual = baseline[valid] - selected
near_residual = np.maximum(residual, 0.0)
near_total = float(np.sum(near_residual))
relative_span = (
float(np.sum(np.square(near_residual))) / near_total
if near_total > 1e-12
else float(np.sqrt(np.mean(np.square(residual))))
)
relative_span = max(relative_span, 1e-6)
depth_signal = np.zeros(depth.shape, dtype=np.float32)
depth_signal[valid] = (reference - selected) / relative_span
baseline_vector = np.asarray(
(
-coefficients[0] / relative_span,
-coefficients[1] / relative_span,
(reference - coefficients[2]) / relative_span,
),
dtype=np.float32,
)
return depth_signal, valid, reference, baseline_vector
def write_depth_analysis(
prediction,
selection_mask,
origin,
source_size,
output_path,
):
import numpy as np
if prediction.depth.shape[0] != 1 or prediction.mask.shape[0] != 1:
raise ValueError("Depth Balloon requires one MoGe-2 frame")
depth_signal, valid, reference, baseline_vector = normalized_moge_depth(
prediction.depth[0],
prediction.mask[0],
selection_mask,
)
crop_size = (selection_mask.shape[1], selection_mask.shape[0])
np.savez_compressed(
output_path,
depth_signal=depth_signal,
baseline_vector=baseline_vector,
valid=valid,
reference_depth=np.asarray(reference, dtype=np.float32),
origin=np.asarray(origin, dtype=np.int32),
crop_size=np.asarray(crop_size, dtype=np.int32),
source_size=np.asarray(source_size, dtype=np.int32),
)
return Path(output_path)
def write_camera_normal_texture(
normal,
valid,
origin,
crop_size,
source_size,
output_path,
):
"""Encode MoGe camera-axis normals into a full-image RGB texture."""
import numpy as np
from PIL import Image
normal = np.asarray(normal, dtype=np.float32)
valid = np.asarray(valid, dtype=bool)
if normal.ndim != 3 or normal.shape[2] != 3:
raise ValueError("Depth Balloon normal data must have three components")
if valid.shape != normal.shape[:2]:
valid = resized_selection_mask(valid, normal.shape[:2])
encoded = np.clip(
(normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32) + 1.0)
* 127.5,
0.0,
255.0,
).astype(np.uint8)
encoded[~valid] = (128, 128, 255)
crop_width, crop_height = (int(value) for value in crop_size)
crop = Image.fromarray(encoded, mode="RGB")
if crop.size != (crop_width, crop_height):
crop = crop.resize(
(crop_width, crop_height),
Image.Resampling.BILINEAR,
)
source_width, source_height = (int(value) for value in source_size)
full = Image.new("RGB", (source_width, source_height), (128, 128, 255))
full.paste(crop, tuple(int(value) for value in origin))
output_path = Path(output_path)
full.save(output_path)
return output_path

View File

@@ -1,26 +0,0 @@
from ..media import copy_frames
from ..models.ben2 import infer_foreground
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing BEN2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
foreground, _load_ms, _inference_ms = infer_foreground(
image,
model_directory,
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
foreground.save(context.directory / f"foreground{suffix}.png")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw BEN2 API foreground")
return {"directory": "."}

View File

@@ -1,21 +0,0 @@
from ..media import copy_frames
from ..models import da3
from ..outputs.debug_da3 import write_prediction
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing DA3 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
prediction = da3.infer(
context.resource("model_manager"),
parameters,
images,
cancel_check=context.check_cancelled,
)
write_prediction(prediction, context.directory / "prediction.npz")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw DA3 API prediction")
return {"directory": "."}

View File

@@ -1,26 +0,0 @@
from ..media import copy_frames
from ..models.moge2 import infer_one
from ..outputs.debug_moge2 import write_prediction
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing MoGe-2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_one(
context.resource("model_manager"),
parameters,
image,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
write_prediction(
prediction,
context.directory / f"prediction{suffix}.npz",
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw MoGe-2 API prediction")
return {"directory": "."}

View File

@@ -0,0 +1,93 @@
from ..media import copy_frames, input_frames
from ..models import da3
from ..models.ben2 import infer_foreground
from ..models.moge2 import infer_one as infer_moge2
from ..models.upscale import infer_one as infer_upscale
from ..outputs.debug_da3 import write_prediction as write_da3_prediction
from ..outputs.debug_moge2 import write_prediction as write_moge2_prediction
def run_da3(context, parameters):
context.progress(0.02, "Preparing DA3 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
prediction = da3.infer(
context.resource("model_manager"),
parameters,
images,
cancel_check=context.check_cancelled,
)
write_da3_prediction(prediction, context.directory / "prediction.npz")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw DA3 API prediction")
return {"directory": "."}
def run_moge2(context, parameters):
context.progress(0.02, "Preparing MoGe-2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_moge2(
context.resource("model_manager"),
parameters,
image,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
write_moge2_prediction(
prediction,
context.directory / f"prediction{suffix}.npz",
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw MoGe-2 API prediction")
return {"directory": "."}
def run_ben2(context, parameters):
context.progress(0.02, "Preparing BEN2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
foreground, _load_ms, _inference_ms = infer_foreground(
image,
model_directory,
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
foreground.save(context.directory / f"foreground{suffix}.png")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw BEN2 API foreground")
return {"directory": "."}
def run_upscale(context, parameters):
context.progress(0.02, "Preparing Real-ESRGAN API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_upscale(
context.resource("model_manager"),
parameters,
image,
cancel_check=context.check_cancelled,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
prediction.save(
context.directory / f"prediction{suffix}.png",
format="PNG",
compress_level=1,
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw Real-ESRGAN API prediction")
return {"directory": "."}

View File

@@ -1,27 +0,0 @@
from ..media import copy_frames
from ..models.upscale import infer_one
from .common import input_frames
def run(context, parameters):
context.progress(0.02, "Preparing Real-ESRGAN API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_one(
context.resource("model_manager"),
parameters,
image,
cancel_check=context.check_cancelled,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
prediction.save(
context.directory / f"prediction{suffix}.png",
format="PNG",
compress_level=1,
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw Real-ESRGAN API prediction")
return {"directory": "."}

View File

@@ -1,51 +0,0 @@
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 analysis_frames, frame_progress_message, input_frames
from .foreground import write_optional
def run(context, parameters):
context.progress(0.02, "Preparing depth input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
with input_frames(context, parameters) as images:
with analysis_frames(context, parameters, images) as model_images:
if adapter is da3:
prediction = adapter.infer(
model_manager,
parameters,
model_images,
progress_callback=lambda index, total: context.progress(
0.15 + 0.65 * index / total,
frame_progress_message("Estimating DA3 depth", index, total),
),
cancel_check=context.check_cancelled,
)
write_da3_depth(prediction, context.directory / "results.npz")
else:
prediction = adapter.infer(
model_manager,
parameters,
model_images,
include_points=False,
progress_callback=lambda index, total: context.progress(
0.15 + 0.65 * index / total,
frame_progress_message("Estimating MoGe-2 depth", index, total),
),
cancel_check=context.check_cancelled,
)
write_moge_depth(context.directory, prediction)
color_paths = write_optional(
context,
parameters,
images,
progress_start=0.84,
progress_span=0.08,
)
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

@@ -1,54 +0,0 @@
from ..models.ben2 import infer_foreground, infer_foreground_frames
def write_optional(
context,
parameters,
images,
*,
progress_start,
progress_span,
):
model_key = parameters.get("remove_background_model", "")
if not model_key:
return []
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(model_key)
device = parameters["device"]
if len(images) == 1:
context.progress(progress_start, "Removing image background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
)
path = context.directory / "image.0001.png"
foreground.convert("RGBA").save(
path,
format="PNG",
compress_level=1,
)
paths = [path]
else:
load_ms, inference_ms = infer_foreground_frames(
images,
model_directory,
device,
context.directory,
context.progress,
progress_start,
progress_span,
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="image",
preserve_input_alpha=True,
)
paths = sorted(context.directory.glob("image.*.png"))
context.log(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)
return paths

View File

@@ -1,21 +0,0 @@
import shutil
from ..media import copy_color_images, prepare_input
def run(context, parameters):
context.directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
context.directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
copy_color_images(images, context.directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
context.progress(1.0, f"Extracted {len(images)} frame(s)")
return {"frames": len(images)}

View File

@@ -1,190 +0,0 @@
import json
import math
import time
from .background import remove_image_background
def polygon_bounds(polygon):
import numpy as np
points = np.asarray(polygon, dtype=np.float64)
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
raise ValueError("The selection does not define a valid image region")
if not np.isfinite(points).all():
raise ValueError("The selection contains invalid image coordinates")
left = int(math.floor(float(points[:, 0].min())))
top = int(math.floor(float(points[:, 1].min())))
right = int(math.ceil(float(points[:, 0].max())))
bottom = int(math.ceil(float(points[:, 1].max())))
if right <= left or bottom <= top:
raise ValueError("The selection does not define a valid image region")
return left, top, right, bottom
def polygon_selection_mask(image_size, polygon, fill_rule="EVEN_ODD"):
"""Rasterize a polygon with an explicit fill rule."""
import numpy as np
width, height = (int(value) for value in image_size)
points = np.asarray(polygon, dtype=np.float64)
if width <= 0 or height <= 0:
raise ValueError("The selection canvas has no readable dimensions")
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
raise ValueError("The selection does not define a valid image region")
if not np.isfinite(points).all():
raise ValueError("The selection contains invalid image coordinates")
if fill_rule not in {"EVEN_ODD", "NONZERO"}:
raise ValueError(f"Unknown polygon fill rule: {fill_rule}")
scanlines = np.arange(height, dtype=np.float64) + 0.5
crossings = [[] for _index in range(height)]
for start, end in zip(points, np.roll(points, -1, axis=0)):
selected = np.flatnonzero(
(start[1] > scanlines) != (end[1] > scanlines)
)
if not len(selected):
continue
crossing_x = start[0] + (
(scanlines[selected] - start[1])
* (end[0] - start[0])
/ (end[1] - start[1])
)
winding = 1 if end[1] > start[1] else -1
for row, value in zip(selected, crossing_x):
crossings[int(row)].append((float(value), winding))
mask = np.zeros((height, width), dtype=bool)
for row, values in enumerate(crossings):
values.sort()
state = 0
index = 0
while index < len(values):
crossing_x = values[index][0]
winding = 0
count = 0
while index < len(values) and values[index][0] == crossing_x:
winding += values[index][1]
count += 1
index += 1
state = (
state + winding
if fill_rule == "NONZERO"
else state ^ (count % 2)
)
if index >= len(values) or state == 0:
continue
following_x = values[index][0]
first = max(0, int(math.ceil(crossing_x - 0.5)))
last = min(width, int(math.floor(following_x - 0.5)) + 1)
mask[row, first:last] = True
return mask
def create_selection_canvas(input_path, polygon, fill_rule="EVEN_ODD"):
"""Create an unbounded selection canvas with transparent black outside."""
import numpy as np
from PIL import Image
with Image.open(input_path) as opened_image:
source = opened_image.convert("RGBA")
left, top, right, bottom = polygon_bounds(polygon)
if right <= 0 or bottom <= 0 or left >= source.width or top >= source.height:
raise ValueError("The selection does not overlap the source image")
canvas = Image.new("RGBA", (right - left, bottom - top), (0, 0, 0, 0))
canvas.paste(source, (-left, -top))
relative = [(float(x) - left, float(y) - top) for x, y in polygon]
selection_mask = polygon_selection_mask(canvas.size, relative, fill_rule)
pixels = np.asarray(canvas, dtype=np.uint8).copy()
pixels[~selection_mask] = 0
return Image.fromarray(pixels, mode="RGBA"), selection_mask, (
left,
top,
right,
bottom,
)
def apply_selection_mask(foreground, selection_mask):
"""Clear pixels outside the source selection after background removal."""
import numpy as np
from PIL import Image
pixels = np.asarray(foreground.convert("RGBA"), dtype=np.uint8).copy()
if pixels.shape[:2] != selection_mask.shape:
pixels = np.asarray(
foreground.convert("RGBA").resize(
(selection_mask.shape[1], selection_mask.shape[0]),
Image.Resampling.BILINEAR,
),
dtype=np.uint8,
)
pixels[~selection_mask] = 0
pixels[pixels[:, :, 3] == 0, :3] = 0
return Image.fromarray(pixels, mode="RGBA")
def crop_to_alpha(image, bounds):
"""Crop an RGBA mask image to its nonzero Alpha bounds."""
import numpy as np
from PIL import Image
pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8)
visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0)
if not len(visible_x):
raise ValueError("The selection contains no visible image pixels")
content_left = int(visible_x.min())
content_top = int(visible_y.min())
content_right = int(visible_x.max()) + 1
content_bottom = int(visible_y.max()) + 1
output = Image.new(
"RGBA",
(
content_right - content_left,
content_bottom - content_top,
),
(0, 0, 0, 0),
)
output.paste(image, (-content_left, -content_top))
left, top, _right, _bottom = bounds
return output, (
left + content_left,
top + content_top,
left + content_right,
top + content_bottom,
)
def run(context, parameters):
"""Refine a gesture mask with BEN2 and return its source-image bounds."""
started = time.perf_counter()
context.progress(0.05, "Preparing image selection")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
canvas, selection_mask, bounds = create_selection_canvas(
parameters["input"], parameters["polygon"]
)
selection_path = output_directory / "selection-input.png"
canvas.save(selection_path, format="PNG", compress_level=1)
model_manager = context.resource("model_manager")
context.progress(0.2, f"Loading BEN2 for {parameters['device'].upper()}")
foreground, _load_ms, _inference_ms = remove_image_background(
selection_path,
model_manager.directory(parameters["model"]),
parameters["device"],
model_manager,
context.check_cancelled,
)
context.check_cancelled()
context.progress(0.9, "Applying selection mask")
output, bounds = crop_to_alpha(
apply_selection_mask(foreground, selection_mask), bounds
)
output_path = output_directory / "image.png"
output.save(output_path, format="PNG", compress_level=1)
bounds_path = output_directory / "bounds.json"
bounds_path.write_text(json.dumps(bounds, separators=(",", ":")), encoding="utf-8")
elapsed_ms = (time.perf_counter() - started) * 1000.0
message = f"Image selection refined in {elapsed_ms:.0f} ms"
context.progress(1.0, message)
return {"image": output_path.name, "bounds": bounds_path.name}

View File

@@ -1,405 +0,0 @@
import json
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from .background import remove_image_background
from .cutout_depth import (
write_camera_normal_texture,
write_depth_analysis,
)
from .common import analysis_frames, input_frames
from .image_mask import (
apply_selection_mask,
create_selection_canvas,
polygon_selection_mask,
)
def _write_field(frame, output_path):
import numpy as np
depth = np.asarray(frame.depth, dtype=np.float32)
points = np.asarray(frame.points, dtype=np.float32)
intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
valid = np.asarray(frame.valid_mask, dtype=bool)
if depth.ndim != 2 or points.shape != (*depth.shape, 3):
raise ValueError("The depth model returned an invalid point field")
valid &= np.isfinite(depth) & (depth > 0.0)
valid &= np.isfinite(points).all(axis=-1)
if not valid.any():
raise ValueError("The depth model returned no valid geometry")
reference_values = depth[valid]
if frame.model_mask is not None:
model_mask = np.asarray(frame.model_mask, dtype=bool)
masked = depth[valid & model_mask]
if len(masked):
reference_values = masked
np.savez_compressed(
output_path,
points=points,
depth=depth,
intrinsics=intrinsics,
valid=valid,
image_size=np.asarray((depth.shape[1], depth.shape[0]), dtype=np.int32),
reference_depth=np.asarray(np.median(reference_values), dtype=np.float32),
)
def _write_masked_color(input_path, polygon, output_path, invert=False):
import numpy as np
from PIL import Image, ImageFilter, ImageOps
with Image.open(input_path) as opened:
source = opened.convert("RGBA")
selection = Image.fromarray(
polygon_selection_mask(source.size, polygon, "NONZERO").astype(np.uint8) * 255,
mode="L",
).filter(ImageFilter.GaussianBlur(radius=0.75))
if invert:
selection = ImageOps.invert(selection)
pixels = np.asarray(source, dtype=np.uint8).copy()
pixels[:, :, 3] = np.asarray(selection, dtype=np.uint8)
pixels[pixels[:, :, 3] == 0, :3] = 0
Image.fromarray(pixels, mode="RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
return output_path
def _write_unbounded_selection(
input_path,
polygon,
output_path,
invert=False,
):
import numpy as np
from PIL import Image, ImageFilter
if invert:
_write_masked_color(input_path, polygon, output_path, invert=True)
with Image.open(input_path) as source:
return output_path, (0, 0, source.width, source.height)
canvas, selection_mask, bounds = create_selection_canvas(
input_path, polygon, 'NONZERO'
)
selection = Image.fromarray(
selection_mask.astype(np.uint8) * 255,
mode='L',
).filter(ImageFilter.GaussianBlur(radius=0.75))
pixels = np.asarray(canvas, dtype=np.uint8).copy()
pixels[:, :, 3] = np.asarray(selection, dtype=np.uint8)
pixels[pixels[:, :, 3] == 0, :3] = 0
Image.fromarray(pixels, mode='RGBA').save(
output_path,
format='PNG',
compress_level=1,
)
return output_path, bounds
def _write_refined_color(context, parameters, input_path, polygon, output_path):
from PIL import Image, ImageChops
canvas, selection_mask, bounds = create_selection_canvas(
input_path, polygon, "NONZERO"
)
selection_path = context.directory / "selection-input.png"
canvas.save(selection_path, format="PNG", compress_level=1)
model_manager = context.resource("model_manager")
foreground, _load_ms, _inference_ms = remove_image_background(
selection_path,
model_manager.directory(parameters["refine_model"]),
parameters["device"],
model_manager,
context.check_cancelled,
)
refined = apply_selection_mask(foreground, selection_mask)
with Image.open(input_path) as opened:
source = opened.convert("RGBA")
if parameters.get("invert"):
mask = Image.new("L", source.size, 0)
mask.paste(refined.getchannel("A"), bounds[:2])
source.putalpha(
ImageChops.multiply(source.getchannel("A"), ImageChops.invert(mask))
)
output = source
else:
output = Image.new("RGBA", source.size, (0, 0, 0, 0))
output.paste(refined, bounds[:2])
output.save(output_path, format="PNG", compress_level=1)
return output_path
def _write_cutout_images(
input_path,
selection_path,
output_directory,
selection_bounds=None,
):
import numpy as np
from PIL import Image
with Image.open(input_path) as opened:
source = opened.convert("RGBA")
with Image.open(selection_path) as opened:
selection = opened.convert("RGBA")
if selection_bounds is None and selection.size != source.size:
raise ValueError("The Cutout selection does not match the source image")
if selection_bounds is None:
selection_bounds = (0, 0, source.width, source.height)
selection_pixels = np.asarray(selection, dtype=np.uint8).copy()
alpha = selection_pixels[:, :, 3]
alpha[alpha <= 1] = 0
selection_pixels[alpha == 0, :3] = 0
selection = Image.fromarray(selection_pixels, mode="RGBA")
visible_y, visible_x = np.nonzero(alpha >= 128)
if not len(visible_x):
raise ValueError("The Cutout selection contains no visible pixels")
bounds = (
int(visible_x.min()),
int(visible_y.min()),
int(visible_x.max()) + 1,
int(visible_y.max()) + 1,
)
content_bounds = bounds
selection_left, selection_top, _selection_right, _selection_bottom = (
selection_bounds
)
bounds = (
selection_left + content_bounds[0],
selection_top + content_bounds[1],
selection_left + content_bounds[2],
selection_top + content_bounds[3],
)
color_canvas = Image.new('RGBA', selection.size, (0, 0, 0, 0))
color_canvas.paste(source, (-selection_left, -selection_top))
color = color_canvas.crop(content_bounds)
color_path = output_directory / "image.0001.png"
color.save(color_path, format="PNG", compress_level=1)
selection_crop = selection.crop(content_bounds)
selection_path = output_directory / "selection.png"
selection_crop.save(selection_path, format="PNG", compress_level=1)
analysis_pixels = np.asarray(color, dtype=np.uint8).copy()
mask_alpha = np.asarray(selection_crop, dtype=np.uint8)[:, :, 3]
analysis_pixels[:, :, 3] = np.minimum(
analysis_pixels[:, :, 3],
mask_alpha,
)
analysis_pixels[mask_alpha < 128] = 0
analysis_path = output_directory / "analysis.png"
Image.fromarray(analysis_pixels, mode="RGBA").save(
analysis_path,
format="PNG",
compress_level=1,
)
bounds_path = output_directory / "bounds.json"
bounds_path.write_text(
json.dumps(bounds, separators=(",", ":")),
encoding="utf-8",
)
return color_path, selection_path, analysis_path, bounds_path
def _write_camera_normal_from_prediction(prediction, color_path, output_path):
import numpy as np
from PIL import Image
normal = np.asarray(prediction.normal[0], dtype=np.float32)
with Image.open(color_path) as source:
alpha = source.convert("RGBA").getchannel("A")
source_size = source.size
valid = (
np.asarray(
alpha.resize((normal.shape[1], normal.shape[0]), Image.Resampling.BILINEAR)
)
> 127
)
return write_camera_normal_texture(
normal,
valid,
(0, 0),
source_size,
source_size,
output_path,
)
def _write_camera_normal(context, parameters, color_path):
normal_parameters = dict(parameters)
normal_parameters.update(
model=parameters["normal_model"],
resolution_level=parameters["normal_resolution_level"],
)
prediction = moge2.infer(
context.resource("model_manager"),
normal_parameters,
[color_path],
include_points=False,
cancel_check=context.check_cancelled,
)
return _write_camera_normal_from_prediction(
prediction,
color_path,
context.directory / "normal.png",
)
def _write_depth_balloon(
context,
parameters,
input_path,
include_normal=False,
):
import numpy as np
from PIL import Image
with Image.open(input_path) as source:
rgba = np.asarray(source.convert("RGBA"), dtype=np.uint8)
source_size = source.size
selection_mask = rgba[:, :, 3] >= 128
prediction = moge2.infer(
context.resource("model_manager"),
parameters,
[input_path],
include_points=False,
cancel_check=context.check_cancelled,
)
depth_path = write_depth_analysis(
prediction,
selection_mask,
(0, 0),
source_size,
context.directory / "depth.npz",
)
normal_path = None
if include_normal:
normal_path = _write_camera_normal_from_prediction(
prediction,
input_path,
context.directory / "normal.png",
)
return depth_path, normal_path
def _write_surface_depth(
context, parameters, color_path, include_normal=False
):
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
if adapter is moge2:
prediction = adapter.infer(
model_manager,
parameters,
[color_path],
include_points=True,
cancel_check=context.check_cancelled,
)
else:
prediction = adapter.infer(
model_manager,
parameters,
[color_path],
cancel_check=context.check_cancelled,
)
context.check_cancelled()
field_path = context.directory / "field.npz"
_write_field(adapter.surface_frame(prediction), field_path)
normal_path = None
if include_normal and adapter is moge2:
normal_path = _write_camera_normal_from_prediction(
prediction, color_path, context.directory / "normal.png"
)
return field_path, normal_path
def run(context, parameters):
context.progress(0.02, "Preparing mesh input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Empty Image to Mesh currently supports one image")
shape = parameters["shape"]
polygon = parameters.get("polygon")
analysis_path = None
result = {}
if polygon:
selection_bounds = None
full_selection_path = context.directory / "selection-full.png"
if parameters.get("refine_selection"):
_write_refined_color(
context,
parameters,
images[0],
polygon,
full_selection_path,
)
else:
_selection_path, selection_bounds = _write_unbounded_selection(
images[0],
polygon,
full_selection_path,
parameters.get("invert", False),
)
color_path, selection_path, analysis_path, bounds_path = (
_write_cutout_images(
images[0],
full_selection_path,
context.directory,
selection_bounds,
)
)
result.update(
selection=selection_path.name,
bounds=bounds_path.name,
)
else:
copy_color_images(images, context.directory)
color_path = min(context.directory.glob("image.0001.*"))
normal_analysis = bool(parameters.get("normal_analysis"))
with analysis_frames(context, parameters, images) as model_images:
if len(model_images) != 1:
raise ValueError("Empty Image to Mesh currently supports one image")
model_path = analysis_path or model_images[0]
if shape in {"CUTOUT", "BALLOON"} and normal_analysis:
context.progress(0.4, "Estimating surface normal")
result["normal"] = _write_camera_normal(
context, parameters, model_path
).name
elif shape == "DEPTH_BALLOON":
context.progress(0.4, "Estimating Balloon depth")
depth_path, normal_path = _write_depth_balloon(
context,
parameters,
model_path,
include_normal=normal_analysis,
)
result["depth"] = depth_path.name
if normal_path is not None:
result["normal"] = normal_path.name
elif shape in {"DEPTH_PLANE", "DEPTH_SURFACE"}:
context.progress(0.4, "Estimating image depth")
field_path, normal_path = _write_surface_depth(
context,
parameters,
model_path,
include_normal=normal_analysis,
)
result["field"] = field_path.name
if (
normal_analysis
and normal_path is None
and shape == "DEPTH_SURFACE"
):
normal_path = _write_camera_normal(
context, parameters, model_path
)
if normal_path is not None:
result["normal"] = normal_path.name
context.check_cancelled()
context.progress(1.0, "Mesh data is ready")
return result

View File

@@ -0,0 +1,255 @@
from ..geometry.normal_texture import write_camera_normal_texture
from ..media import (
DEFAULT_MAX_AI_INPUT_SIZE,
analysis_frames,
copy_color_images,
input_frames,
limit_image,
)
from ..models import geometry_adapter, moge2
from ..selection.cutout_images import refine_selection, write_images, write_selection
def _limited_model_input(source_path, output_path, max_input_size):
from PIL import Image
with Image.open(source_path) as opened:
limited = limit_image(opened, max_input_size)
if limited is opened:
return source_path
limited.save(output_path, format="PNG", compress_level=1)
return output_path
def _write_field(frame, output_path, calibration_valid=None):
import numpy as np
depth = np.asarray(frame.depth, dtype=np.float32)
points = np.asarray(frame.points, dtype=np.float32)
intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
valid = np.asarray(frame.valid_mask, dtype=bool)
if depth.ndim != 2 or points.shape != (*depth.shape, 3):
raise ValueError("The depth model returned an invalid point field")
valid &= np.isfinite(depth) & (depth > 0.0)
valid &= np.isfinite(points).all(axis=-1)
if not valid.any():
raise ValueError("The depth model returned no valid geometry")
if calibration_valid is None:
calibration_valid = valid.copy()
else:
calibration_valid = np.asarray(calibration_valid, dtype=bool)
if calibration_valid.shape != valid.shape:
raise ValueError("The depth calibration mask dimensions do not match")
calibration_valid &= valid
if not calibration_valid.any():
raise ValueError("The depth field contains no valid calibration pixels")
reference_values = depth[calibration_valid]
if frame.model_mask is not None:
model_mask = np.asarray(frame.model_mask, dtype=bool)
masked = depth[calibration_valid & model_mask]
if len(masked):
reference_values = masked
np.savez_compressed(
output_path,
points=points,
depth=depth,
intrinsics=intrinsics,
valid=valid,
calibration_valid=calibration_valid,
image_size=np.asarray((depth.shape[1], depth.shape[0]), dtype=np.int32),
reference_depth=np.asarray(np.median(reference_values), dtype=np.float32),
)
def _selection_calibration_mask(frame, selection_path):
import numpy as np
from PIL import Image
valid = np.asarray(frame.valid_mask, dtype=bool)
with Image.open(selection_path) as opened:
selection = opened.convert("RGBA").getchannel("A")
selection = selection.resize(
(valid.shape[1], valid.shape[0]),
Image.Resampling.NEAREST,
)
selected = np.asarray(selection, dtype=np.uint8) >= 26
calibration_valid = valid & selected
if not calibration_valid.any():
raise ValueError("The depth field contains no valid selected pixels")
return calibration_valid
def _write_camera_normal_from_prediction(prediction, color_path, output_path):
import numpy as np
from PIL import Image
normal = np.asarray(prediction.normal[0], dtype=np.float32)
with Image.open(color_path) as source:
alpha = source.convert("RGBA").getchannel("A")
source_size = source.size
valid = (
np.asarray(
alpha.resize((normal.shape[1], normal.shape[0]), Image.Resampling.BILINEAR)
)
> 127
)
return write_camera_normal_texture(
normal,
valid,
(0, 0),
source_size,
source_size,
output_path,
)
def _write_camera_normal(context, parameters, model_path, color_path=None):
prediction = moge2.infer(
context.resource("model_manager"),
parameters,
[model_path],
include_points=False,
cancel_check=context.check_cancelled,
)
return _write_camera_normal_from_prediction(
prediction,
color_path or model_path,
context.directory / "normal.png",
)
def _write_depth_field(
context,
parameters,
model_path,
normal_color_path=None,
include_normal=False,
selection_path=None,
):
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
if adapter is moge2:
prediction = adapter.infer(
model_manager,
parameters,
[model_path],
include_points=True,
cancel_check=context.check_cancelled,
)
else:
prediction = adapter.infer(
model_manager,
parameters,
[model_path],
cancel_check=context.check_cancelled,
)
context.check_cancelled()
field_path = context.directory / "field.npz"
frame = adapter.geometry_frame(prediction)
calibration_valid = frame.valid_mask
if selection_path is not None:
calibration_valid = _selection_calibration_mask(frame, selection_path)
_write_field(frame, field_path, calibration_valid)
normal_path = None
if include_normal and adapter is moge2:
normal_path = _write_camera_normal_from_prediction(
prediction,
normal_color_path or model_path,
context.directory / "normal.png",
)
return field_path, normal_path
def run(context, parameters):
context.progress(0.02, "Preparing geometry input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Empty Image to Mesh currently supports one image")
shape = parameters["shape"]
polygon = parameters.get("polygon")
analysis_path = None
selection_path = None
result = {}
if polygon:
selection_bounds = None
full_selection_path = context.directory / "selection-full.png"
if parameters.get("refine_selection"):
refine_selection(
context,
parameters,
images[0],
polygon,
full_selection_path,
)
else:
_selection_path, selection_bounds = write_selection(
images[0],
polygon,
full_selection_path,
parameters.get("invert", False),
)
color_path, selection_path, analysis_path, bounds_path = (
write_images(
images[0],
full_selection_path,
context.directory,
selection_bounds,
)
)
result.update(
selection=selection_path.name,
bounds=bounds_path.name,
)
else:
copy_color_images(images, context.directory)
normal_analysis = bool(parameters.get("normal_analysis"))
with analysis_frames(context, parameters, images) as model_images:
if len(model_images) != 1:
raise ValueError("Empty Image to Mesh currently supports one image")
model_path = analysis_path or model_images[0]
limited_model_path = model_path
if normal_analysis or shape in {
"DEPTH_PLANE",
"DEPTH_BALLOON",
"DEPTH_SURFACE",
}:
limited_model_path = _limited_model_input(
model_path,
context.directory / "model-input.png",
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
)
if shape in {"CUTOUT", "BALLOON"} and normal_analysis:
context.progress(0.4, "Estimating surface normal")
result["normal"] = _write_camera_normal(
context, parameters, limited_model_path, model_path
).name
elif shape in {"DEPTH_PLANE", "DEPTH_BALLOON", "DEPTH_SURFACE"}:
context.progress(0.4, "Estimating image depth")
field_path, normal_path = _write_depth_field(
context,
parameters,
limited_model_path,
normal_color_path=model_path,
include_normal=normal_analysis,
selection_path=selection_path,
)
result["field"] = field_path.name
if (
normal_analysis
and normal_path is None
and shape == "DEPTH_SURFACE"
):
normal_path = _write_camera_normal(
context, parameters, limited_model_path, model_path
)
if normal_path is not None:
result["normal"] = normal_path.name
if limited_model_path != model_path:
limited_model_path.unlink(missing_ok=True)
context.check_cancelled()
context.progress(1.0, "Geometry data is ready")
return result

View File

@@ -0,0 +1,46 @@
import json
import time
from ..models.ben2 import infer_foreground
from ..selection.polygon_mask import (
apply_selection_mask,
create_selection_canvas,
crop_to_alpha,
)
def run(context, parameters):
"""Refine a gesture mask with BEN2 and return its source-image bounds."""
started = time.perf_counter()
context.progress(0.05, "Preparing image selection")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
canvas, selection_mask, bounds = create_selection_canvas(
parameters["input"], parameters["polygon"]
)
selection_path = output_directory / "selection-input.png"
canvas.save(selection_path, format="PNG", compress_level=1)
model_manager = context.resource("model_manager")
context.progress(0.2, f"Loading BEN2 for {parameters['device'].upper()}")
foreground, _load_ms, _inference_ms = infer_foreground(
selection_path,
model_manager.directory(parameters["model"]),
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
)
context.check_cancelled()
context.progress(0.9, "Applying selection mask")
output, bounds = crop_to_alpha(
apply_selection_mask(foreground, selection_mask), bounds
)
output_path = output_directory / "image.png"
output.save(output_path, format="PNG", compress_level=1)
bounds_path = output_directory / "bounds.json"
bounds_path.write_text(json.dumps(bounds, separators=(",", ":")), encoding="utf-8")
elapsed_ms = (time.perf_counter() - started) * 1000.0
message = f"Image selection refined in {elapsed_ms:.0f} ms"
context.progress(1.0, message)
return {"image": output_path.name, "bounds": bounds_path.name}

View File

@@ -1,24 +1,9 @@
import shutil
import time
from ..media import prepare_background_input
from ..media import input_frames
from ..models.ben2 import infer_foreground, infer_foreground_frames
def remove_image_background(
input_path, model_directory, device, model_manager, cancel_check
):
"""Remove one image background with the shared BEN2 configuration."""
return infer_foreground(
input_path,
model_directory,
device,
model_cache=model_manager,
cancel_check=cancel_check,
preserve_input_alpha=True,
)
def run(context, parameters):
"""Run BEN2 and write transparent foreground images into the Job directory."""
total_started = time.perf_counter()
@@ -31,25 +16,17 @@ def run(context, parameters):
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
video_frames = parameters.get("video_frames", 16)
start_frame = parameters.get("start_frame", 1)
images, temporary_directory = prepare_background_input(
parameters["input"],
video_frames,
output_directory,
start_frame=start_frame,
cancel_check=context.check_cancelled,
)
try:
with input_frames(context, parameters) as images:
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
context.progress(0.35, "Removing background")
foreground, load_ms, inference_ms = remove_image_background(
foreground, load_ms, inference_ms = infer_foreground(
images[0],
model_directory,
device,
model_manager,
context.check_cancelled,
model_cache=model_manager,
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
@@ -75,9 +52,6 @@ def run(context, parameters):
output_prefix="image",
preserve_input_alpha=True,
)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
context.check_cancelled()
total_ms = (time.perf_counter() - total_started) * 1000.0

View File

@@ -1,71 +0,0 @@
from ..geometry import build_surface, foreground_mesh_mask
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from ..outputs.surface import write_normal_texture, write_surface_archive
from .common import analysis_frames, frame_progress_message, input_frames
from .foreground import write_optional
def run(context, parameters):
context.progress(0.02, "Preparing surface input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Convert to Surface currently supports one image")
context.progress(0.12, "Estimating surface geometry")
with analysis_frames(context, parameters, images) as model_images:
if len(model_images) != 1:
raise ValueError("Convert to Surface currently supports one image")
if adapter is moge2:
prediction = adapter.infer(
model_manager,
parameters,
model_images,
include_points=True,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
frame_progress_message("Estimating MoGe-2 geometry", index, total),
),
cancel_check=context.check_cancelled,
)
else:
prediction = adapter.infer(
model_manager,
parameters,
model_images,
progress_callback=lambda index, total: context.progress(
0.12 + 0.58 * index / total,
frame_progress_message("Estimating DA3 geometry", index, total),
),
cancel_check=context.check_cancelled,
)
context.check_cancelled()
frame = adapter.surface_frame(prediction)
color_paths = write_optional(
context,
parameters,
images,
progress_start=0.72,
progress_span=0.08,
)
foreground_mask = None
if color_paths:
foreground_mask = foreground_mesh_mask(
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")
if not color_paths:
copy_color_images(images, context.directory)
if adapter is moge2:
write_normal_texture(
prediction.normal[0],
context.directory / "normal.png",
)
context.check_cancelled()
context.progress(1.0, "Surface is ready")
return {"surface": "mesh.npz"}

View File

@@ -1,11 +1,18 @@
from ..media import DEFAULT_MAX_AI_INPUT_SIZE, input_frames
from ..models.upscale import infer_one
from .common import frame_progress_message, input_frames
def frame_progress_message(message, index, total):
if total > 1:
return f"{message} ({index + 1}/{total})"
return message
def run(context, parameters):
context.log(
f"Upscale model: {parameters['model']}; "
f"maximum output size: {parameters.get('max_output_size', 4096)} px"
f"maximum AI input size: "
f"{parameters.get('max_input_size', DEFAULT_MAX_AI_INPUT_SIZE)} px"
)
context.progress(0.02, "Preparing upscale input")
context.directory.mkdir(parents=True, exist_ok=True)