Add new node asset scripts for geometry and shader operations
- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters. - Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections. - Added `validate.py` to validate the integrity of the newly created node groups and their inputs. - Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies. - Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
This commit is contained in:
@@ -1,37 +0,0 @@
|
||||
import shutil
|
||||
from contextlib import contextmanager
|
||||
|
||||
from ..media import prepare_input
|
||||
|
||||
|
||||
def frame_progress_message(message, index, total):
|
||||
if total > 1:
|
||||
return f"{message} ({index + 1}/{total})"
|
||||
return message
|
||||
|
||||
|
||||
@contextmanager
|
||||
def input_frames(context, parameters):
|
||||
images, temporary_directory = prepare_input(
|
||||
parameters["input"],
|
||||
parameters.get("video_frames", 16),
|
||||
context.directory,
|
||||
start_frame=parameters.get("start_frame", 1),
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
try:
|
||||
yield images
|
||||
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
|
||||
@@ -1,176 +0,0 @@
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PLANE_VARIANCE_FLOOR_RATIO = 0.025
|
||||
|
||||
|
||||
def resized_selection_mask(selection_mask, shape):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
selection_mask = np.asarray(selection_mask, dtype=bool)
|
||||
if selection_mask.shape == tuple(shape):
|
||||
return selection_mask
|
||||
resized = Image.fromarray(selection_mask.astype(np.uint8) * 255, mode="L")
|
||||
resized = resized.resize(
|
||||
(int(shape[1]), int(shape[0])),
|
||||
Image.Resampling.NEAREST,
|
||||
)
|
||||
return np.asarray(resized, dtype=np.uint8) > 0
|
||||
|
||||
|
||||
def _stable_weighted_plane(coordinates, values, weights):
|
||||
import numpy as np
|
||||
|
||||
total = float(np.sum(weights))
|
||||
if total <= 1e-12:
|
||||
return np.zeros(3, dtype=np.float64)
|
||||
center = np.sum(coordinates * weights[:, None], axis=0) / total
|
||||
value_center = float(np.sum(values * weights) / total)
|
||||
centered = coordinates - center
|
||||
covariance = (centered * weights[:, None]).T @ centered / total
|
||||
cross = centered.T @ (weights * (values - value_center)) / total
|
||||
eigenvalues, eigenvectors = np.linalg.eigh(covariance)
|
||||
largest = max(float(eigenvalues[-1]), 0.0)
|
||||
if largest <= 1e-12:
|
||||
slopes = np.zeros(2, dtype=np.float64)
|
||||
else:
|
||||
supported = np.maximum(eigenvalues, largest * PLANE_VARIANCE_FLOOR_RATIO)
|
||||
slopes = eigenvectors @ ((eigenvectors.T @ cross) / supported)
|
||||
intercept = value_center - float(center @ slopes)
|
||||
return np.asarray((*slopes, intercept), dtype=np.float64)
|
||||
|
||||
|
||||
def normalized_moge_depth(depth, model_mask, selection_mask):
|
||||
"""Normalize MoGe depth while preserving its adjustable fitted base plane."""
|
||||
import numpy as np
|
||||
|
||||
depth = np.asarray(depth, dtype=np.float32)
|
||||
model_mask = np.asarray(model_mask, dtype=bool)
|
||||
selection_mask = resized_selection_mask(selection_mask, depth.shape)
|
||||
if model_mask.shape != depth.shape:
|
||||
raise ValueError("MoGe-2 depth and mask dimensions do not match")
|
||||
valid = selection_mask & model_mask & np.isfinite(depth) & (depth > 0.0)
|
||||
if not valid.any():
|
||||
raise ValueError("MoGe-2 found no valid depth inside the lasso")
|
||||
|
||||
height, width = depth.shape
|
||||
y, x = np.indices(depth.shape, dtype=np.float64)
|
||||
x = x / max(width - 1, 1) * 2.0 - 1.0
|
||||
y = y / max(height - 1, 1) * 2.0 - 1.0
|
||||
coordinates = np.stack((x[valid], y[valid]), axis=1)
|
||||
selected = depth[valid].astype(np.float64)
|
||||
weights = np.ones_like(selected)
|
||||
coefficients = np.zeros(3, dtype=np.float64)
|
||||
for _iteration in range(4):
|
||||
coefficients = _stable_weighted_plane(
|
||||
coordinates,
|
||||
selected,
|
||||
weights,
|
||||
)
|
||||
residual = (
|
||||
selected
|
||||
- coordinates @ coefficients[:2]
|
||||
- coefficients[2]
|
||||
)
|
||||
centered = residual - np.median(residual)
|
||||
scale = max(float(np.median(np.abs(centered))) * 1.4826, 1e-6)
|
||||
weights = np.minimum(
|
||||
1.0,
|
||||
(2.5 * scale) / np.maximum(np.abs(centered), 1e-12),
|
||||
)
|
||||
|
||||
baseline = coefficients[0] * x + coefficients[1] * y + coefficients[2]
|
||||
reference = float(np.median(baseline[valid]))
|
||||
residual = baseline[valid] - selected
|
||||
near_residual = np.maximum(residual, 0.0)
|
||||
near_total = float(np.sum(near_residual))
|
||||
relative_span = (
|
||||
float(np.sum(np.square(near_residual))) / near_total
|
||||
if near_total > 1e-12
|
||||
else float(np.sqrt(np.mean(np.square(residual))))
|
||||
)
|
||||
relative_span = max(relative_span, 1e-6)
|
||||
depth_signal = np.zeros(depth.shape, dtype=np.float32)
|
||||
depth_signal[valid] = (reference - selected) / relative_span
|
||||
baseline_vector = np.asarray(
|
||||
(
|
||||
-coefficients[0] / relative_span,
|
||||
-coefficients[1] / relative_span,
|
||||
(reference - coefficients[2]) / relative_span,
|
||||
),
|
||||
dtype=np.float32,
|
||||
)
|
||||
return depth_signal, valid, reference, baseline_vector
|
||||
|
||||
|
||||
def write_depth_analysis(
|
||||
prediction,
|
||||
selection_mask,
|
||||
origin,
|
||||
source_size,
|
||||
output_path,
|
||||
):
|
||||
import numpy as np
|
||||
|
||||
if prediction.depth.shape[0] != 1 or prediction.mask.shape[0] != 1:
|
||||
raise ValueError("Depth Balloon requires one MoGe-2 frame")
|
||||
depth_signal, valid, reference, baseline_vector = normalized_moge_depth(
|
||||
prediction.depth[0],
|
||||
prediction.mask[0],
|
||||
selection_mask,
|
||||
)
|
||||
crop_size = (selection_mask.shape[1], selection_mask.shape[0])
|
||||
np.savez_compressed(
|
||||
output_path,
|
||||
depth_signal=depth_signal,
|
||||
baseline_vector=baseline_vector,
|
||||
valid=valid,
|
||||
reference_depth=np.asarray(reference, dtype=np.float32),
|
||||
origin=np.asarray(origin, dtype=np.int32),
|
||||
crop_size=np.asarray(crop_size, dtype=np.int32),
|
||||
source_size=np.asarray(source_size, dtype=np.int32),
|
||||
)
|
||||
return Path(output_path)
|
||||
|
||||
|
||||
def write_camera_normal_texture(
|
||||
normal,
|
||||
valid,
|
||||
origin,
|
||||
crop_size,
|
||||
source_size,
|
||||
output_path,
|
||||
):
|
||||
"""Encode MoGe camera-axis normals into a full-image RGB texture."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
normal = np.asarray(normal, dtype=np.float32)
|
||||
valid = np.asarray(valid, dtype=bool)
|
||||
if normal.ndim != 3 or normal.shape[2] != 3:
|
||||
raise ValueError("Depth Balloon normal data must have three components")
|
||||
if valid.shape != normal.shape[:2]:
|
||||
valid = resized_selection_mask(valid, normal.shape[:2])
|
||||
encoded = np.clip(
|
||||
(normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32) + 1.0)
|
||||
* 127.5,
|
||||
0.0,
|
||||
255.0,
|
||||
).astype(np.uint8)
|
||||
encoded[~valid] = (128, 128, 255)
|
||||
crop_width, crop_height = (int(value) for value in crop_size)
|
||||
crop = Image.fromarray(encoded, mode="RGB")
|
||||
if crop.size != (crop_width, crop_height):
|
||||
crop = crop.resize(
|
||||
(crop_width, crop_height),
|
||||
Image.Resampling.BILINEAR,
|
||||
)
|
||||
source_width, source_height = (int(value) for value in source_size)
|
||||
full = Image.new("RGB", (source_width, source_height), (128, 128, 255))
|
||||
full.paste(crop, tuple(int(value) for value in origin))
|
||||
output_path = Path(output_path)
|
||||
full.save(output_path)
|
||||
return output_path
|
||||
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
from ..media import copy_frames
|
||||
from ..models.ben2 import infer_foreground
|
||||
from .common import input_frames
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing BEN2 API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
model_manager = context.resource("model_manager")
|
||||
model_directory = model_manager.directory(parameters["model"])
|
||||
with input_frames(context, parameters) as images:
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
foreground, _load_ms, _inference_ms = infer_foreground(
|
||||
image,
|
||||
model_directory,
|
||||
parameters["device"],
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
|
||||
foreground.save(context.directory / f"foreground{suffix}.png")
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw BEN2 API foreground")
|
||||
return {"directory": "."}
|
||||
@@ -1,21 +0,0 @@
|
||||
from ..media import copy_frames
|
||||
from ..models import da3
|
||||
from ..outputs.debug_da3 import write_prediction
|
||||
from .common import input_frames
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing DA3 API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
prediction = da3.infer(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
images,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
write_prediction(prediction, context.directory / "prediction.npz")
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw DA3 API prediction")
|
||||
return {"directory": "."}
|
||||
@@ -1,26 +0,0 @@
|
||||
from ..media import copy_frames
|
||||
from ..models.moge2 import infer_one
|
||||
from ..outputs.debug_moge2 import write_prediction
|
||||
from .common import input_frames
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing MoGe-2 API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
prediction = infer_one(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
image,
|
||||
)
|
||||
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
|
||||
write_prediction(
|
||||
prediction,
|
||||
context.directory / f"prediction{suffix}.npz",
|
||||
)
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw MoGe-2 API prediction")
|
||||
return {"directory": "."}
|
||||
93
src/anyimage/server/jobs/debug_predictions.py
Normal file
93
src/anyimage/server/jobs/debug_predictions.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from ..media import copy_frames, input_frames
|
||||
from ..models import da3
|
||||
from ..models.ben2 import infer_foreground
|
||||
from ..models.moge2 import infer_one as infer_moge2
|
||||
from ..models.upscale import infer_one as infer_upscale
|
||||
from ..outputs.debug_da3 import write_prediction as write_da3_prediction
|
||||
from ..outputs.debug_moge2 import write_prediction as write_moge2_prediction
|
||||
|
||||
|
||||
def run_da3(context, parameters):
|
||||
context.progress(0.02, "Preparing DA3 API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
prediction = da3.infer(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
images,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
write_da3_prediction(prediction, context.directory / "prediction.npz")
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw DA3 API prediction")
|
||||
return {"directory": "."}
|
||||
|
||||
|
||||
def run_moge2(context, parameters):
|
||||
context.progress(0.02, "Preparing MoGe-2 API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
prediction = infer_moge2(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
image,
|
||||
)
|
||||
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
|
||||
write_moge2_prediction(
|
||||
prediction,
|
||||
context.directory / f"prediction{suffix}.npz",
|
||||
)
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw MoGe-2 API prediction")
|
||||
return {"directory": "."}
|
||||
|
||||
|
||||
def run_ben2(context, parameters):
|
||||
context.progress(0.02, "Preparing BEN2 API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
model_manager = context.resource("model_manager")
|
||||
model_directory = model_manager.directory(parameters["model"])
|
||||
with input_frames(context, parameters) as images:
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
foreground, _load_ms, _inference_ms = infer_foreground(
|
||||
image,
|
||||
model_directory,
|
||||
parameters["device"],
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
|
||||
foreground.save(context.directory / f"foreground{suffix}.png")
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw BEN2 API foreground")
|
||||
return {"directory": "."}
|
||||
|
||||
|
||||
def run_upscale(context, parameters):
|
||||
context.progress(0.02, "Preparing Real-ESRGAN API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
prediction = infer_upscale(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
image,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
|
||||
prediction.save(
|
||||
context.directory / f"prediction{suffix}.png",
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw Real-ESRGAN API prediction")
|
||||
return {"directory": "."}
|
||||
@@ -1,27 +0,0 @@
|
||||
from ..media import copy_frames
|
||||
from ..models.upscale import infer_one
|
||||
from .common import input_frames
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing Real-ESRGAN API input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
for index, image in enumerate(images):
|
||||
context.check_cancelled()
|
||||
prediction = infer_one(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
image,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
|
||||
prediction.save(
|
||||
context.directory / f"prediction{suffix}.png",
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
copy_frames(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Saved raw Real-ESRGAN API prediction")
|
||||
return {"directory": "."}
|
||||
@@ -1,51 +0,0 @@
|
||||
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 analysis_frames, frame_progress_message, input_frames
|
||||
from .foreground import write_optional
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing depth input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
model_manager = context.resource("model_manager")
|
||||
adapter = geometry_adapter(parameters["model"])
|
||||
with input_frames(context, parameters) as images:
|
||||
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,
|
||||
images,
|
||||
progress_start=0.84,
|
||||
progress_span=0.08,
|
||||
)
|
||||
if not color_paths:
|
||||
copy_color_images(images, context.directory)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, f"Exported depth for {len(images)} frame(s)")
|
||||
return {"depth": "results.npz"}
|
||||
@@ -1,54 +0,0 @@
|
||||
from ..models.ben2 import infer_foreground, infer_foreground_frames
|
||||
|
||||
|
||||
def write_optional(
|
||||
context,
|
||||
parameters,
|
||||
images,
|
||||
*,
|
||||
progress_start,
|
||||
progress_span,
|
||||
):
|
||||
model_key = parameters.get("remove_background_model", "")
|
||||
if not model_key:
|
||||
return []
|
||||
model_manager = context.resource("model_manager")
|
||||
model_directory = model_manager.directory(model_key)
|
||||
device = parameters["device"]
|
||||
if len(images) == 1:
|
||||
context.progress(progress_start, "Removing image background")
|
||||
foreground, load_ms, inference_ms = infer_foreground(
|
||||
images[0],
|
||||
model_directory,
|
||||
device,
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
path = context.directory / "image.0001.png"
|
||||
foreground.convert("RGBA").save(
|
||||
path,
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
paths = [path]
|
||||
else:
|
||||
load_ms, inference_ms = infer_foreground_frames(
|
||||
images,
|
||||
model_directory,
|
||||
device,
|
||||
context.directory,
|
||||
context.progress,
|
||||
progress_start,
|
||||
progress_span,
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
output_prefix="image",
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
paths = sorted(context.directory.glob("image.*.png"))
|
||||
context.log(
|
||||
f"BEN2 inference {inference_ms:.0f} ms; "
|
||||
f"model load {load_ms:.0f} ms"
|
||||
)
|
||||
return paths
|
||||
@@ -1,21 +0,0 @@
|
||||
import shutil
|
||||
|
||||
from ..media import copy_color_images, prepare_input
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
images, temporary_directory = prepare_input(
|
||||
parameters["input"],
|
||||
parameters.get("video_frames", 16),
|
||||
context.directory,
|
||||
start_frame=parameters.get("start_frame", 1),
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
try:
|
||||
copy_color_images(images, context.directory)
|
||||
finally:
|
||||
if temporary_directory is not None:
|
||||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||
context.progress(1.0, f"Extracted {len(images)} frame(s)")
|
||||
return {"frames": len(images)}
|
||||
@@ -1,190 +0,0 @@
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
|
||||
from .background import remove_image_background
|
||||
|
||||
|
||||
def polygon_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 selection does not define a valid image region")
|
||||
if not np.isfinite(points).all():
|
||||
raise ValueError("The selection 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 selection does not define a valid image region")
|
||||
return left, top, right, bottom
|
||||
|
||||
|
||||
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
|
||||
|
||||
with Image.open(input_path) as opened_image:
|
||||
source = opened_image.convert("RGBA")
|
||||
left, top, right, bottom = polygon_bounds(polygon)
|
||||
if right <= 0 or bottom <= 0 or left >= source.width or top >= source.height:
|
||||
raise ValueError("The selection 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_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, (
|
||||
left,
|
||||
top,
|
||||
right,
|
||||
bottom,
|
||||
)
|
||||
|
||||
|
||||
def apply_selection_mask(foreground, selection_mask):
|
||||
"""Clear pixels outside the source selection after background removal."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
pixels = np.asarray(foreground.convert("RGBA"), dtype=np.uint8).copy()
|
||||
if pixels.shape[:2] != selection_mask.shape:
|
||||
pixels = np.asarray(
|
||||
foreground.convert("RGBA").resize(
|
||||
(selection_mask.shape[1], selection_mask.shape[0]),
|
||||
Image.Resampling.BILINEAR,
|
||||
),
|
||||
dtype=np.uint8,
|
||||
)
|
||||
pixels[~selection_mask] = 0
|
||||
pixels[pixels[:, :, 3] == 0, :3] = 0
|
||||
return Image.fromarray(pixels, mode="RGBA")
|
||||
|
||||
|
||||
def crop_to_alpha(image, bounds):
|
||||
"""Crop an RGBA mask image to its nonzero Alpha bounds."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
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 selection 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,
|
||||
content_bottom - content_top,
|
||||
),
|
||||
(0, 0, 0, 0),
|
||||
)
|
||||
output.paste(image, (-content_left, -content_top))
|
||||
left, top, _right, _bottom = bounds
|
||||
return output, (
|
||||
left + content_left,
|
||||
top + content_top,
|
||||
left + content_right,
|
||||
top + content_bottom,
|
||||
)
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
"""Refine a gesture mask with BEN2 and return its source-image bounds."""
|
||||
started = time.perf_counter()
|
||||
context.progress(0.05, "Preparing image selection")
|
||||
context.check_cancelled()
|
||||
output_directory = context.directory
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
canvas, selection_mask, bounds = create_selection_canvas(
|
||||
parameters["input"], parameters["polygon"]
|
||||
)
|
||||
selection_path = output_directory / "selection-input.png"
|
||||
canvas.save(selection_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 = remove_image_background(
|
||||
selection_path,
|
||||
model_manager.directory(parameters["model"]),
|
||||
parameters["device"],
|
||||
model_manager,
|
||||
context.check_cancelled,
|
||||
)
|
||||
context.check_cancelled()
|
||||
context.progress(0.9, "Applying selection mask")
|
||||
output, bounds = crop_to_alpha(
|
||||
apply_selection_mask(foreground, selection_mask), bounds
|
||||
)
|
||||
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 selection refined in {elapsed_ms:.0f} ms"
|
||||
context.progress(1.0, message)
|
||||
return {"image": output_path.name, "bounds": bounds_path.name}
|
||||
@@ -1,405 +0,0 @@
|
||||
import json
|
||||
|
||||
from ..media import copy_color_images
|
||||
from ..models import geometry_adapter, moge2
|
||||
from .background import remove_image_background
|
||||
from .cutout_depth 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_depth_balloon(
|
||||
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_depth_balloon(
|
||||
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
|
||||
255
src/anyimage/server/jobs/prepare_geometry_data.py
Normal file
255
src/anyimage/server/jobs/prepare_geometry_data.py
Normal file
@@ -0,0 +1,255 @@
|
||||
from ..geometry.normal_texture import write_camera_normal_texture
|
||||
from ..media import (
|
||||
DEFAULT_MAX_AI_INPUT_SIZE,
|
||||
analysis_frames,
|
||||
copy_color_images,
|
||||
input_frames,
|
||||
limit_image,
|
||||
)
|
||||
from ..models import geometry_adapter, moge2
|
||||
from ..selection.cutout_images import refine_selection, write_images, write_selection
|
||||
|
||||
|
||||
def _limited_model_input(source_path, output_path, max_input_size):
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(source_path) as opened:
|
||||
limited = limit_image(opened, max_input_size)
|
||||
if limited is opened:
|
||||
return source_path
|
||||
limited.save(output_path, format="PNG", compress_level=1)
|
||||
return output_path
|
||||
|
||||
|
||||
def _write_field(frame, output_path, calibration_valid=None):
|
||||
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")
|
||||
if calibration_valid is None:
|
||||
calibration_valid = valid.copy()
|
||||
else:
|
||||
calibration_valid = np.asarray(calibration_valid, dtype=bool)
|
||||
if calibration_valid.shape != valid.shape:
|
||||
raise ValueError("The depth calibration mask dimensions do not match")
|
||||
calibration_valid &= valid
|
||||
if not calibration_valid.any():
|
||||
raise ValueError("The depth field contains no valid calibration pixels")
|
||||
reference_values = depth[calibration_valid]
|
||||
if frame.model_mask is not None:
|
||||
model_mask = np.asarray(frame.model_mask, dtype=bool)
|
||||
masked = depth[calibration_valid & model_mask]
|
||||
if len(masked):
|
||||
reference_values = masked
|
||||
np.savez_compressed(
|
||||
output_path,
|
||||
points=points,
|
||||
depth=depth,
|
||||
intrinsics=intrinsics,
|
||||
valid=valid,
|
||||
calibration_valid=calibration_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 _selection_calibration_mask(frame, selection_path):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
valid = np.asarray(frame.valid_mask, dtype=bool)
|
||||
with Image.open(selection_path) as opened:
|
||||
selection = opened.convert("RGBA").getchannel("A")
|
||||
selection = selection.resize(
|
||||
(valid.shape[1], valid.shape[0]),
|
||||
Image.Resampling.NEAREST,
|
||||
)
|
||||
selected = np.asarray(selection, dtype=np.uint8) >= 26
|
||||
calibration_valid = valid & selected
|
||||
if not calibration_valid.any():
|
||||
raise ValueError("The depth field contains no valid selected pixels")
|
||||
return calibration_valid
|
||||
|
||||
|
||||
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, model_path, color_path=None):
|
||||
prediction = moge2.infer(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
[model_path],
|
||||
include_points=False,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
return _write_camera_normal_from_prediction(
|
||||
prediction,
|
||||
color_path or model_path,
|
||||
context.directory / "normal.png",
|
||||
)
|
||||
|
||||
|
||||
def _write_depth_field(
|
||||
context,
|
||||
parameters,
|
||||
model_path,
|
||||
normal_color_path=None,
|
||||
include_normal=False,
|
||||
selection_path=None,
|
||||
):
|
||||
model_manager = context.resource("model_manager")
|
||||
adapter = geometry_adapter(parameters["model"])
|
||||
if adapter is moge2:
|
||||
prediction = adapter.infer(
|
||||
model_manager,
|
||||
parameters,
|
||||
[model_path],
|
||||
include_points=True,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
else:
|
||||
prediction = adapter.infer(
|
||||
model_manager,
|
||||
parameters,
|
||||
[model_path],
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
context.check_cancelled()
|
||||
field_path = context.directory / "field.npz"
|
||||
frame = adapter.geometry_frame(prediction)
|
||||
calibration_valid = frame.valid_mask
|
||||
if selection_path is not None:
|
||||
calibration_valid = _selection_calibration_mask(frame, selection_path)
|
||||
_write_field(frame, field_path, calibration_valid)
|
||||
normal_path = None
|
||||
if include_normal and adapter is moge2:
|
||||
normal_path = _write_camera_normal_from_prediction(
|
||||
prediction,
|
||||
normal_color_path or model_path,
|
||||
context.directory / "normal.png",
|
||||
)
|
||||
return field_path, normal_path
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing geometry 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
|
||||
selection_path = None
|
||||
result = {}
|
||||
if polygon:
|
||||
selection_bounds = None
|
||||
full_selection_path = context.directory / "selection-full.png"
|
||||
if parameters.get("refine_selection"):
|
||||
refine_selection(
|
||||
context,
|
||||
parameters,
|
||||
images[0],
|
||||
polygon,
|
||||
full_selection_path,
|
||||
)
|
||||
else:
|
||||
_selection_path, selection_bounds = write_selection(
|
||||
images[0],
|
||||
polygon,
|
||||
full_selection_path,
|
||||
parameters.get("invert", False),
|
||||
)
|
||||
color_path, selection_path, analysis_path, bounds_path = (
|
||||
write_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)
|
||||
|
||||
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]
|
||||
limited_model_path = model_path
|
||||
if normal_analysis or shape in {
|
||||
"DEPTH_PLANE",
|
||||
"DEPTH_BALLOON",
|
||||
"DEPTH_SURFACE",
|
||||
}:
|
||||
limited_model_path = _limited_model_input(
|
||||
model_path,
|
||||
context.directory / "model-input.png",
|
||||
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
|
||||
)
|
||||
if shape in {"CUTOUT", "BALLOON"} and normal_analysis:
|
||||
context.progress(0.4, "Estimating surface normal")
|
||||
result["normal"] = _write_camera_normal(
|
||||
context, parameters, limited_model_path, model_path
|
||||
).name
|
||||
elif shape in {"DEPTH_PLANE", "DEPTH_BALLOON", "DEPTH_SURFACE"}:
|
||||
context.progress(0.4, "Estimating image depth")
|
||||
field_path, normal_path = _write_depth_field(
|
||||
context,
|
||||
parameters,
|
||||
limited_model_path,
|
||||
normal_color_path=model_path,
|
||||
include_normal=normal_analysis,
|
||||
selection_path=selection_path,
|
||||
)
|
||||
result["field"] = field_path.name
|
||||
if (
|
||||
normal_analysis
|
||||
and normal_path is None
|
||||
and shape == "DEPTH_SURFACE"
|
||||
):
|
||||
normal_path = _write_camera_normal(
|
||||
context, parameters, limited_model_path, model_path
|
||||
)
|
||||
if normal_path is not None:
|
||||
result["normal"] = normal_path.name
|
||||
|
||||
if limited_model_path != model_path:
|
||||
limited_model_path.unlink(missing_ok=True)
|
||||
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Geometry data is ready")
|
||||
return result
|
||||
46
src/anyimage/server/jobs/refine_image_mask.py
Normal file
46
src/anyimage/server/jobs/refine_image_mask.py
Normal file
@@ -0,0 +1,46 @@
|
||||
import json
|
||||
import time
|
||||
|
||||
from ..models.ben2 import infer_foreground
|
||||
from ..selection.polygon_mask import (
|
||||
apply_selection_mask,
|
||||
create_selection_canvas,
|
||||
crop_to_alpha,
|
||||
)
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
"""Refine a gesture mask with BEN2 and return its source-image bounds."""
|
||||
started = time.perf_counter()
|
||||
context.progress(0.05, "Preparing image selection")
|
||||
context.check_cancelled()
|
||||
output_directory = context.directory
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
canvas, selection_mask, bounds = create_selection_canvas(
|
||||
parameters["input"], parameters["polygon"]
|
||||
)
|
||||
selection_path = output_directory / "selection-input.png"
|
||||
canvas.save(selection_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(
|
||||
selection_path,
|
||||
model_manager.directory(parameters["model"]),
|
||||
parameters["device"],
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
context.check_cancelled()
|
||||
context.progress(0.9, "Applying selection mask")
|
||||
output, bounds = crop_to_alpha(
|
||||
apply_selection_mask(foreground, selection_mask), bounds
|
||||
)
|
||||
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 selection refined in {elapsed_ms:.0f} ms"
|
||||
context.progress(1.0, message)
|
||||
return {"image": output_path.name, "bounds": bounds_path.name}
|
||||
@@ -1,24 +1,9 @@
|
||||
import shutil
|
||||
import time
|
||||
|
||||
from ..media import prepare_background_input
|
||||
from ..media import input_frames
|
||||
from ..models.ben2 import infer_foreground, infer_foreground_frames
|
||||
|
||||
|
||||
def remove_image_background(
|
||||
input_path, model_directory, device, model_manager, cancel_check
|
||||
):
|
||||
"""Remove one image background with the shared BEN2 configuration."""
|
||||
return infer_foreground(
|
||||
input_path,
|
||||
model_directory,
|
||||
device,
|
||||
model_cache=model_manager,
|
||||
cancel_check=cancel_check,
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
"""Run BEN2 and write transparent foreground images into the Job directory."""
|
||||
total_started = time.perf_counter()
|
||||
@@ -31,25 +16,17 @@ def run(context, parameters):
|
||||
model_manager = context.resource("model_manager")
|
||||
model_directory = model_manager.directory(parameters["model"])
|
||||
|
||||
video_frames = parameters.get("video_frames", 16)
|
||||
start_frame = parameters.get("start_frame", 1)
|
||||
images, temporary_directory = prepare_background_input(
|
||||
parameters["input"],
|
||||
video_frames,
|
||||
output_directory,
|
||||
start_frame=start_frame,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
try:
|
||||
with input_frames(context, parameters) as images:
|
||||
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
|
||||
if len(images) == 1:
|
||||
context.progress(0.35, "Removing background")
|
||||
foreground, load_ms, inference_ms = remove_image_background(
|
||||
foreground, load_ms, inference_ms = infer_foreground(
|
||||
images[0],
|
||||
model_directory,
|
||||
device,
|
||||
model_manager,
|
||||
context.check_cancelled,
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
context.check_cancelled()
|
||||
context.progress(0.92, "Writing transparent PNG")
|
||||
@@ -75,9 +52,6 @@ def run(context, parameters):
|
||||
output_prefix="image",
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
finally:
|
||||
if temporary_directory is not None:
|
||||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||
|
||||
context.check_cancelled()
|
||||
total_ms = (time.perf_counter() - total_started) * 1000.0
|
||||
@@ -1,71 +0,0 @@
|
||||
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 analysis_frames, frame_progress_message, input_frames
|
||||
from .foreground import write_optional
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.02, "Preparing surface input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
model_manager = context.resource("model_manager")
|
||||
adapter = geometry_adapter(parameters["model"])
|
||||
with input_frames(context, parameters) as images:
|
||||
if len(images) != 1:
|
||||
raise ValueError("Convert to Surface currently supports one image")
|
||||
context.progress(0.12, "Estimating surface geometry")
|
||||
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(
|
||||
context,
|
||||
parameters,
|
||||
images,
|
||||
progress_start=0.72,
|
||||
progress_span=0.08,
|
||||
)
|
||||
foreground_mask = None
|
||||
if color_paths:
|
||||
foreground_mask = foreground_mesh_mask(
|
||||
color_paths[0],
|
||||
frame.depth.shape,
|
||||
)
|
||||
context.progress(0.82, "Building surface")
|
||||
mesh = build_surface(frame, foreground_mask=foreground_mask)
|
||||
context.progress(0.92, "Writing surface")
|
||||
write_surface_archive(mesh, context.directory / "mesh.npz")
|
||||
if not color_paths:
|
||||
copy_color_images(images, context.directory)
|
||||
if adapter is moge2:
|
||||
write_normal_texture(
|
||||
prediction.normal[0],
|
||||
context.directory / "normal.png",
|
||||
)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Surface is ready")
|
||||
return {"surface": "mesh.npz"}
|
||||
@@ -1,11 +1,18 @@
|
||||
from ..media import DEFAULT_MAX_AI_INPUT_SIZE, input_frames
|
||||
from ..models.upscale import infer_one
|
||||
from .common import frame_progress_message, input_frames
|
||||
|
||||
|
||||
def frame_progress_message(message, index, total):
|
||||
if total > 1:
|
||||
return f"{message} ({index + 1}/{total})"
|
||||
return message
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.log(
|
||||
f"Upscale model: {parameters['model']}; "
|
||||
f"maximum output size: {parameters.get('max_output_size', 4096)} px"
|
||||
f"maximum AI input size: "
|
||||
f"{parameters.get('max_input_size', DEFAULT_MAX_AI_INPUT_SIZE)} px"
|
||||
)
|
||||
context.progress(0.02, "Preparing upscale input")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
Reference in New Issue
Block a user