Add tests for clipboard image functionality and mesh plane asset creation

- Implemented a comprehensive test suite for the clipboard image module, covering various functionalities including operator registration, context handling, and image processing.
- Created a new mesh plane asset with customizable parameters and integrated it into the geometry node system.
- Developed a validation script to ensure the correct behavior of the mesh plane asset, including input handling and output verification.
This commit is contained in:
2026-08-22 09:49:00 +08:00
parent af5603045c
commit 988332548e
49 changed files with 4419 additions and 291 deletions

View File

@@ -10,6 +10,7 @@ from .jobs.debug_ben2 import run as _debug_ben2
from .jobs.debug_da3 import run as _debug_da3
from .jobs.debug_moge2 import run as _debug_moge2
from .jobs.frames import run as _extract_frames
from .jobs.image_lasso import run as _analyze_image_lasso
from .model_manager import ModelManager, ModelSessionCache
@@ -38,6 +39,11 @@ def analyze_balloon(context, parameters):
return _analyze_balloon(context, parameters)
@server.job("analyze-image-lasso")
def analyze_image_lasso(context, parameters):
return _analyze_image_lasso(context, parameters)
@server.job("remove-background")
def remove_background(context, parameters):
return _remove_background(context, parameters)

View File

@@ -36,6 +36,7 @@ def run(context, parameters):
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
@@ -59,6 +60,7 @@ def run(context, parameters):
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
preserve_input_alpha=True,
)
finally:
if temporary_directory is not None:

View File

@@ -23,6 +23,7 @@ def write_optional(
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
)
path = context.directory / "color.0001.png"
foreground.convert("RGBA").save(
@@ -43,6 +44,7 @@ def write_optional(
model_cache=model_manager,
cancel_check=context.check_cancelled,
output_prefix="color",
preserve_input_alpha=True,
)
paths = sorted(context.directory.glob("color.*.png"))
print(

View File

@@ -0,0 +1,151 @@
import json
import math
import time
from ..models.ben2 import infer_foreground
def lasso_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 lasso does not define a valid image region")
if not np.isfinite(points).all():
raise ValueError("The lasso 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 lasso does not define a valid image region")
return left, top, right, bottom
def create_lasso_canvas(input_path, polygon):
"""Create an unbounded lasso canvas with transparent black outside."""
import numpy as np
from PIL import Image, ImageDraw
with Image.open(input_path) as opened_image:
source = opened_image.convert("RGBA")
left, top, right, bottom = lasso_bounds(polygon)
if right <= 0 or bottom <= 0 or left >= source.width or top >= source.height:
raise ValueError("The lasso 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 = Image.new("L", canvas.size, 0)
ImageDraw.Draw(selection).polygon(relative, fill=255)
selection_mask = np.asarray(selection, dtype=np.uint8) > 0
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_contour_alpha(canvas, selection_mask, foreground):
"""Keep source RGB and combine source, lasso, and BEN2 alpha."""
import numpy as np
from PIL import Image
source = np.asarray(canvas.convert("RGBA"), dtype=np.uint8).copy()
alpha = np.asarray(foreground.convert("RGBA").getchannel("A"), dtype=np.uint8)
if alpha.shape != selection_mask.shape:
alpha = np.asarray(
Image.fromarray(alpha, mode="L").resize(
canvas.size, Image.Resampling.BILINEAR
),
dtype=np.uint8,
)
source[:, :, 3] = np.where(
selection_mask,
np.minimum(source[:, :, 3], alpha),
0,
).astype(np.uint8)
source[source[:, :, 3] == 0, :3] = 0
return Image.fromarray(source, mode="RGBA")
def crop_to_alpha(image, bounds, padding=0):
"""Crop an RGBA image to nonzero Alpha and expand it by pixel padding."""
import numpy as np
from PIL import Image
if padding < 0:
raise ValueError("Padding must not be negative")
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 lasso 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 + padding * 2,
content_bottom - content_top + padding * 2,
),
(0, 0, 0, 0),
)
output.paste(image, (padding - content_left, padding - content_top))
left, top, _right, _bottom = bounds
return output, (
left + content_left - padding,
top + content_top - padding,
left + content_right + padding,
top + content_bottom + padding,
)
def run(context, parameters):
"""Use BEN2 to refine an Image Lasso without changing its canvas bounds."""
started = time.perf_counter()
context.progress(0.05, "Preparing image lasso")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
canvas, selection_mask, bounds = create_lasso_canvas(
parameters["input"], parameters["polygon"]
)
lasso_path = output_directory / "lasso-input.png"
import numpy as np
from PIL import Image
analysis_pixels = np.asarray(canvas, dtype=np.uint8).copy()
analysis_pixels[~selection_mask] = 0
Image.fromarray(analysis_pixels, mode="RGBA").save(
lasso_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(
lasso_path,
model_manager.directory(parameters["model"]),
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(0.9, "Applying contour alpha")
output, bounds = crop_to_alpha(
apply_contour_alpha(canvas, selection_mask, foreground),
bounds,
int(parameters.get("padding", 0)),
)
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 contour finished in {elapsed_ms:.0f} ms"
context.progress(1.0, message)
return {"image": output_path.name, "bounds": bounds_path.name}

View File

@@ -10,8 +10,9 @@ def infer_foreground(
device,
model_cache=None,
cancel_check=None,
preserve_input_alpha=False,
):
from PIL import Image
from PIL import Image, ImageChops
input_path = Path(input_path)
check_cancelled(cancel_check)
@@ -30,10 +31,29 @@ def infer_foreground(
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
source_rgba = opened_image.convert("RGBA")
if preserve_input_alpha:
image = Image.new("RGB", source_rgba.size, (0, 0, 0))
image.paste(
source_rgba.convert("RGB"),
(0, 0),
source_rgba.getchannel("A"),
)
else:
image = source_rgba.convert("RGB")
inference_started = time.perf_counter()
foreground = infer(model, image)
inference_ms = (time.perf_counter() - inference_started) * 1000.0
if preserve_input_alpha:
foreground_alpha = foreground.convert("RGBA").getchannel("A")
source_alpha = source_rgba.getchannel("A")
if source_alpha.size != foreground_alpha.size:
source_alpha = source_alpha.resize(
foreground_alpha.size,
Image.Resampling.BILINEAR,
)
source_rgba.putalpha(ImageChops.multiply(source_alpha, foreground_alpha))
foreground = source_rgba
check_cancelled(cancel_check)
return foreground, load_ms, inference_ms
@@ -49,6 +69,7 @@ def infer_foreground_frames(
model_cache=None,
cancel_check=None,
output_prefix="foreground",
preserve_input_alpha=False,
):
"""Run BEN2 on every frame and write numbered transparent PNG files."""
total_load_ms = 0.0
@@ -65,6 +86,7 @@ def infer_foreground_frames(
device,
model_cache=model_cache,
cancel_check=cancel_check,
preserve_input_alpha=preserve_input_alpha,
)
output_path = output_directory / f"{output_prefix}.{index + 1:04d}.png"
foreground.convert("RGBA").save(