Add comprehensive tests for cutout and mesh functionalities
- 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:
303
src/anyimage/server/jobs/mesh.py
Normal file
303
src/anyimage/server/jobs/mesh.py
Normal 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
|
||||
Reference in New Issue
Block a user