nan
2b79a342b0
- 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.
304 lines
10 KiB
Python
304 lines
10 KiB
Python
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
|