icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add comprehensive tests for cutout and mesh functionalities

Browse Source 
- Introduced `test_cutout.py` to validate cutout pie menu behavior, image lasso options, and mesh detail calculations.
- Added `test_mesh.py` to ensure correct tangent normal rasterization, plane source topology, and depth surface evaluations.
- Created `test_mesh_blender.py` to verify material node groups, tangent baking, and local cutout building from Blender images.
- Implemented tests for handling alpha regions and ensuring geometry nodes maintain expected behavior.
This commit is contained in:
马南
2026-08-26 10:42:29 +08:00
parent 0cf04a13c6
commit 2b79a342b0
46 changed files with 5554 additions and 645 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/anyimage/server/jobs/__init__.py
View File

@@ -1,5 +1 @@
from .depth import run as generate_depth
from .surface import run as generate_surface
__all__ = ("generate_depth", "generate_surface")
"""Job implementations registered by :mod:`anyimage.server.app`."""

20
src/anyimage/server/jobs/balloon.py
View File

@@ -5,15 +5,18 @@ from pathlib import Path
from ..models import moge2
from ..models.ben2 import infer_foreground
from .image_mask import polygon_selection_mask
ALPHA_THRESHOLD = 0.5
def create_lasso_input(input_path, polygon, output_path):
def create_lasso_input(
input_path, polygon, output_path, fill_rule="EVEN_ODD"
):
"""Crop a lasso and clear every RGBA channel outside its polygon."""
import numpy as np
from PIL import Image, ImageDraw
from PIL import Image
points = np.asarray(polygon, dtype=np.float64)
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
@@ -33,12 +36,7 @@ def create_lasso_input(input_path, polygon, output_path):
relative = points - np.asarray((left, top), dtype=np.float64)
cropped = source.crop((left, top, right, bottom))
selection = Image.new("L", cropped.size, 0)
ImageDraw.Draw(selection).polygon(
[tuple(point) for point in relative],
fill=255,
)
selection_mask = np.asarray(selection, dtype=np.uint8) > 0
selection_mask = polygon_selection_mask(cropped.size, relative, fill_rule)
pixels = np.asarray(cropped, dtype=np.uint8).copy()
pixels[~selection_mask] = 0
masked = Image.fromarray(pixels, mode="RGBA")
@@ -285,7 +283,7 @@ def write_depth_analysis(
return Path(output_path)
def write_balloon_normal_texture(
def write_camera_normal_texture(
normal,
valid,
origin,
@@ -293,7 +291,7 @@ def write_balloon_normal_texture(
source_size,
output_path,
):
"""Write MoGe camera normals as a full-image tangent-space texture."""
"""Encode MoGe camera-axis normals into a full-image RGB texture."""
import numpy as np
from PIL import Image
@@ -412,7 +410,7 @@ def run(context, parameters):
import numpy as np
with np.load(depth_path, allow_pickle=False) as data:
normal_path = write_balloon_normal_texture(
normal_path = write_camera_normal_texture(
prediction.normal[0],
data["valid"],
origin,

11
src/anyimage/server/jobs/common.py
View File

@@ -24,3 +24,14 @@ def input_frames(context, parameters):
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

53
src/anyimage/server/jobs/depth.py
View File

@@ -1,7 +1,7 @@
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 frame_progress_message, input_frames
from .common import analysis_frames, frame_progress_message, input_frames
from .foreground import write_optional
@@ -11,31 +11,32 @@ def run(context, parameters):
model_manager = context.resource("model_manager")
adapter = geometry_adapter(parameters["model"])
with input_frames(context, parameters) as images:
if adapter is da3:
prediction = adapter.infer(
model_manager,
parameters,
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,
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)
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,

66
src/anyimage/server/jobs/image_mask.py
View File

@@ -22,10 +22,68 @@ def polygon_bounds(polygon):
return left, top, right, bottom
def create_selection_canvas(input_path, polygon):
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, ImageDraw
from PIL import Image
with Image.open(input_path) as opened_image:
source = opened_image.convert("RGBA")
@@ -35,9 +93,7 @@ def create_selection_canvas(input_path, polygon):
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
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, (

303
src/anyimage/server/jobs/mesh.py Normal file
View File

@@ -0,0 +1,303 @@
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from .background import remove_image_background
from .balloon 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, ImageChops, 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)
alpha = ImageChops.multiply(source.getchannel("A"), selection)
pixels = np.asarray(source, dtype=np.uint8).copy()
pixels[:, :, 3] = np.asarray(alpha, 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_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_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_balloon_depth(
context,
parameters,
input_path,
polygon,
selection_size=None,
include_normal=False,
):
from PIL import Image
with Image.open(input_path) as source:
source_size = source.size
selection_size = selection_size or source_size
scale_x = source_size[0] / selection_size[0]
scale_y = source_size[1] / selection_size[1]
analysis_polygon = tuple(
(point[0] * scale_x, point[1] * scale_y) for point in polygon
)
selection_mask = polygon_selection_mask(
source_size,
analysis_polygon,
"NONZERO",
)
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")
if polygon:
if parameters.get("refine_selection"):
color_path = _write_refined_color(
context,
parameters,
images[0],
polygon,
context.directory / "image.0001.png",
)
else:
color_path = _write_masked_color(
images[0],
polygon,
context.directory / "image.0001.png",
parameters.get("invert", False),
)
else:
copy_color_images(images, context.directory)
color_path = min(context.directory.glob("image.0001.*"))
result = {}
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")
analysis_path = 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, analysis_path
).name
elif shape == "DEPTH_BALLOON":
from PIL import Image
with Image.open(images[0]) as selection_source:
selection_size = selection_source.size
context.progress(0.4, "Estimating Balloon depth")
depth_path, normal_path = _write_balloon_depth(
context,
parameters,
analysis_path,
polygon,
selection_size,
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,
analysis_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, analysis_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

51
src/anyimage/server/jobs/surface.py
View File

@@ -2,7 +2,7 @@ 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 frame_progress_message, input_frames
from .common import analysis_frames, frame_progress_message, input_frames
from .foreground import write_optional
def run(context, parameters):
@@ -14,29 +14,32 @@ def run(context, parameters):
if len(images) != 1:
raise ValueError("Convert to Surface currently supports one image")
context.progress(0.12, "Estimating surface geometry")
if adapter is moge2:
prediction = adapter.infer(
model_manager,
parameters,
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,
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,
)
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(