nan
6c8358f92f
- Updated the naming convention for upscaled images to include underscores and file extensions. - Renamed properties related to cutout pixel size to point spacing for clarity. - Adjusted alpha region calculations to use a half-alpha threshold for better accuracy. - Enhanced the handling of foreground mesh masks to consider half-alpha values. - Introduced new functions for unbounded selections and cutout images, ensuring proper handling of source and selection images. - Improved test coverage for cutout functionalities, including bounds mapping and alpha preservation. - Updated tests to reflect changes in naming conventions and functionality adjustments.
406 lines
14 KiB
Python
406 lines
14 KiB
Python
import json
|
|
|
|
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, 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)
|
|
pixels = np.asarray(source, dtype=np.uint8).copy()
|
|
pixels[:, :, 3] = np.asarray(selection, 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_unbounded_selection(
|
|
input_path,
|
|
polygon,
|
|
output_path,
|
|
invert=False,
|
|
):
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
if invert:
|
|
_write_masked_color(input_path, polygon, output_path, invert=True)
|
|
with Image.open(input_path) as source:
|
|
return output_path, (0, 0, source.width, source.height)
|
|
canvas, selection_mask, bounds = create_selection_canvas(
|
|
input_path, polygon, 'NONZERO'
|
|
)
|
|
selection = Image.fromarray(
|
|
selection_mask.astype(np.uint8) * 255,
|
|
mode='L',
|
|
).filter(ImageFilter.GaussianBlur(radius=0.75))
|
|
pixels = np.asarray(canvas, dtype=np.uint8).copy()
|
|
pixels[:, :, 3] = np.asarray(selection, 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, bounds
|
|
|
|
|
|
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_cutout_images(
|
|
input_path,
|
|
selection_path,
|
|
output_directory,
|
|
selection_bounds=None,
|
|
):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
with Image.open(input_path) as opened:
|
|
source = opened.convert("RGBA")
|
|
with Image.open(selection_path) as opened:
|
|
selection = opened.convert("RGBA")
|
|
if selection_bounds is None and selection.size != source.size:
|
|
raise ValueError("The Cutout selection does not match the source image")
|
|
if selection_bounds is None:
|
|
selection_bounds = (0, 0, source.width, source.height)
|
|
selection_pixels = np.asarray(selection, dtype=np.uint8).copy()
|
|
alpha = selection_pixels[:, :, 3]
|
|
alpha[alpha <= 1] = 0
|
|
selection_pixels[alpha == 0, :3] = 0
|
|
selection = Image.fromarray(selection_pixels, mode="RGBA")
|
|
visible_y, visible_x = np.nonzero(alpha >= 128)
|
|
if not len(visible_x):
|
|
raise ValueError("The Cutout selection contains no visible pixels")
|
|
bounds = (
|
|
int(visible_x.min()),
|
|
int(visible_y.min()),
|
|
int(visible_x.max()) + 1,
|
|
int(visible_y.max()) + 1,
|
|
)
|
|
content_bounds = bounds
|
|
selection_left, selection_top, _selection_right, _selection_bottom = (
|
|
selection_bounds
|
|
)
|
|
bounds = (
|
|
selection_left + content_bounds[0],
|
|
selection_top + content_bounds[1],
|
|
selection_left + content_bounds[2],
|
|
selection_top + content_bounds[3],
|
|
)
|
|
color_canvas = Image.new('RGBA', selection.size, (0, 0, 0, 0))
|
|
color_canvas.paste(source, (-selection_left, -selection_top))
|
|
color = color_canvas.crop(content_bounds)
|
|
color_path = output_directory / "image.0001.png"
|
|
color.save(color_path, format="PNG", compress_level=1)
|
|
|
|
selection_crop = selection.crop(content_bounds)
|
|
selection_path = output_directory / "selection.png"
|
|
selection_crop.save(selection_path, format="PNG", compress_level=1)
|
|
|
|
analysis_pixels = np.asarray(color, dtype=np.uint8).copy()
|
|
mask_alpha = np.asarray(selection_crop, dtype=np.uint8)[:, :, 3]
|
|
analysis_pixels[:, :, 3] = np.minimum(
|
|
analysis_pixels[:, :, 3],
|
|
mask_alpha,
|
|
)
|
|
analysis_pixels[mask_alpha < 128] = 0
|
|
analysis_path = output_directory / "analysis.png"
|
|
Image.fromarray(analysis_pixels, mode="RGBA").save(
|
|
analysis_path,
|
|
format="PNG",
|
|
compress_level=1,
|
|
)
|
|
bounds_path = output_directory / "bounds.json"
|
|
bounds_path.write_text(
|
|
json.dumps(bounds, separators=(",", ":")),
|
|
encoding="utf-8",
|
|
)
|
|
return color_path, selection_path, analysis_path, bounds_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,
|
|
include_normal=False,
|
|
):
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
with Image.open(input_path) as source:
|
|
rgba = np.asarray(source.convert("RGBA"), dtype=np.uint8)
|
|
source_size = source.size
|
|
selection_mask = rgba[:, :, 3] >= 128
|
|
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")
|
|
analysis_path = None
|
|
result = {}
|
|
if polygon:
|
|
selection_bounds = None
|
|
full_selection_path = context.directory / "selection-full.png"
|
|
if parameters.get("refine_selection"):
|
|
_write_refined_color(
|
|
context,
|
|
parameters,
|
|
images[0],
|
|
polygon,
|
|
full_selection_path,
|
|
)
|
|
else:
|
|
_selection_path, selection_bounds = _write_unbounded_selection(
|
|
images[0],
|
|
polygon,
|
|
full_selection_path,
|
|
parameters.get("invert", False),
|
|
)
|
|
color_path, selection_path, analysis_path, bounds_path = (
|
|
_write_cutout_images(
|
|
images[0],
|
|
full_selection_path,
|
|
context.directory,
|
|
selection_bounds,
|
|
)
|
|
)
|
|
result.update(
|
|
selection=selection_path.name,
|
|
bounds=bounds_path.name,
|
|
)
|
|
else:
|
|
copy_color_images(images, context.directory)
|
|
color_path = min(context.directory.glob("image.0001.*"))
|
|
|
|
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")
|
|
model_path = analysis_path or 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, model_path
|
|
).name
|
|
elif shape == "DEPTH_BALLOON":
|
|
context.progress(0.4, "Estimating Balloon depth")
|
|
depth_path, normal_path = _write_balloon_depth(
|
|
context,
|
|
parameters,
|
|
model_path,
|
|
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,
|
|
model_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, model_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
|