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:
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
151
src/anyimage/server/jobs/image_lasso.py
Normal file
151
src/anyimage/server/jobs/image_lasso.py
Normal 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}
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user