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

Add cutout profile and teddy balloon height calculations

Browse Source 
- Implemented cotangent Poisson field and Poisson balloon profile in cutout_profile.py.
- Added teddy balloon heights calculation and related functions in cutout_teddy.py.
- Introduced interaction utilities for image editing in interaction.py.
- Developed depth normalization and analysis functions in cutout_depth.py.
- Created a validation script for cutout assets in validate_cutout_asset.py.
This commit is contained in:
马南
2026-08-27 22:46:57 +08:00
parent 6c8358f92f
commit 7d39fe8c5e
40 changed files with 3058 additions and 3896 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,426 +0,0 @@
import json
import math
import time
from pathlib import Path
from ..models import moge2
from ..models.ben2 import infer_foreground
from .image_mask import polygon_selection_mask
ALPHA_THRESHOLD = 0.5
def create_lasso_input(
input_path, polygon, output_path, fill_rule="EVEN_ODD"
):
"""Crop a lasso and clear every RGBA channel outside its polygon."""
import numpy as np
from PIL import Image
points = np.asarray(polygon, dtype=np.float64)
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
raise ValueError("The lasso does not define a valid image region")
if not np.isfinite(points).all():
raise ValueError("The lasso contains invalid image coordinates")
with Image.open(input_path) as opened_image:
source = opened_image.convert("RGBA")
width, height = source.size
left = max(0, int(math.floor(float(points[:, 0].min()))))
top = max(0, int(math.floor(float(points[:, 1].min()))))
right = min(width, int(math.ceil(float(points[:, 0].max()))))
bottom = min(height, int(math.ceil(float(points[:, 1].max()))))
if right <= left or bottom <= top:
raise ValueError("The lasso does not overlap the source image")
relative = points - np.asarray((left, top), dtype=np.float64)
cropped = source.crop((left, top, right, bottom))
selection_mask = polygon_selection_mask(cropped.size, relative, fill_rule)
pixels = np.asarray(cropped, dtype=np.uint8).copy()
pixels[~selection_mask] = 0
masked = Image.fromarray(pixels, mode="RGBA")
output_path = Path(output_path)
masked.save(output_path, format="PNG", compress_level=1)
return output_path, selection_mask, (left, top), (width, height)
def largest_component(mask):
"""Return the largest four-connected component of a boolean image."""
import numpy as np
mask = np.asarray(mask, dtype=bool)
if mask.ndim != 2:
raise ValueError("The BEN2 Alpha mask must be two-dimensional")
height, width = mask.shape
visited = np.zeros_like(mask)
largest = []
for start_y, start_x in np.argwhere(mask):
if visited[start_y, start_x]:
continue
visited[start_y, start_x] = True
pending = [(int(start_x), int(start_y))]
component = []
while pending:
x, y = pending.pop()
component.append((x, y))
for next_x, next_y in (
(x - 1, y),
(x + 1, y),
(x, y - 1),
(x, y + 1),
):
if (
0 <= next_x < width
and 0 <= next_y < height
and mask[next_y, next_x]
and not visited[next_y, next_x]
):
visited[next_y, next_x] = True
pending.append((next_x, next_y))
if len(component) > len(largest):
largest = component
if not largest:
raise ValueError("BEN2 found no foreground inside the lasso")
result = np.zeros_like(mask)
component_x, component_y = zip(*largest)
result[component_y, component_x] = True
return result
def _choose_boundary_edge(start, direction, candidates):
right = (-direction[1], direction[0])
left = (direction[1], -direction[0])
back = (-direction[0], -direction[1])
preference = {right: 0, direction: 1, left: 2, back: 3}
return min(
candidates,
key=lambda end: preference.get(
(end[0] - start[0], end[1] - start[1]),
4,
),
)
def _simplify_axis_aligned_contour(points):
simplified = []
for index, current in enumerate(points):
previous = points[index - 1]
following = points[(index + 1) % len(points)]
first = (current[0] - previous[0], current[1] - previous[1])
second = (following[0] - current[0], following[1] - current[1])
if first[0] * second[1] - first[1] * second[0] != 0:
simplified.append(current)
return simplified
def trace_component_contour(component):
"""Trace the outer pixel-edge loop of one four-connected component."""
import numpy as np
component = np.asarray(component, dtype=bool)
height, width = component.shape
edges = {}
def add_edge(start, end):
edges.setdefault(start, []).append(end)
for y, x in np.argwhere(component):
x = int(x)
y = int(y)
if y == 0 or not component[y - 1, x]:
add_edge((x, y), (x + 1, y))
if x == width - 1 or not component[y, x + 1]:
add_edge((x + 1, y), (x + 1, y + 1))
if y == height - 1 or not component[y + 1, x]:
add_edge((x + 1, y + 1), (x, y + 1))
if x == 0 or not component[y, x - 1]:
add_edge((x, y + 1), (x, y))
if not edges:
raise ValueError("BEN2 foreground has no usable boundary")
start = min(edges, key=lambda point: (point[1], point[0]))
current = start
direction = (1, 0)
contour = []
edge_count = sum(len(values) for values in edges.values())
for _index in range(edge_count + 1):
contour.append(current)
candidates = edges.get(current, ())
if not candidates:
raise ValueError("BEN2 foreground boundary is open")
following = _choose_boundary_edge(current, direction, candidates)
direction = (following[0] - current[0], following[1] - current[1])
current = following
if current == start:
break
else:
raise ValueError("BEN2 foreground boundary could not be closed")
contour = _simplify_axis_aligned_contour(contour)
if len(contour) < 3:
raise ValueError("BEN2 foreground is too small for a Balloon")
return contour
def contour_from_alpha(alpha, selection_mask, origin=(0, 0)):
"""Build a lasso contour from BEN2 Alpha values strictly above 0.5."""
import numpy as np
alpha = np.asarray(alpha, dtype=np.float32)
selection_mask = np.asarray(selection_mask, dtype=bool)
if alpha.shape != selection_mask.shape:
raise ValueError("BEN2 Alpha does not match the lasso input")
foreground = (alpha / 255.0 > ALPHA_THRESHOLD) & selection_mask
component = largest_component(foreground)
contour = trace_component_contour(component)
offset_x, offset_y = origin
return [[x + offset_x, y + offset_y] for x, y in contour]
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 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
design = np.stack((x[valid], y[valid], np.ones(valid.sum())), 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):
weighted_design = design * np.sqrt(weights)[:, None]
weighted_depth = selected * np.sqrt(weights)
coefficients = np.linalg.lstsq(
weighted_design,
weighted_depth,
rcond=None,
)[0]
residual = selected - design @ coefficients
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("Balloon Depth 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("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
def run(context, parameters):
"""Analyze a Balloon lasso with optional BEN2 contour and MoGe-2 depth."""
total_started = time.perf_counter()
contour_analysis = bool(parameters.get("contour_analysis"))
depth_analysis = bool(parameters.get("depth_analysis"))
if not contour_analysis and not depth_analysis:
raise ValueError("Select Contour or Depth")
context.progress(0.02, "Preparing lasso input")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
input_path, selection_mask, origin, source_size = create_lasso_input(
parameters["input"],
parameters["polygon"],
output_directory / "lasso-input.png",
)
context.check_cancelled()
device = parameters["device"]
model_manager = context.resource("model_manager")
result = {}
if contour_analysis:
model_directory = model_manager.directory(parameters["contour_model"])
context.progress(0.15, f"Loading BEN2 for {device.upper()}")
foreground, _load_ms, _inference_ms = infer_foreground(
input_path,
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(
0.55 if depth_analysis else 0.9,
"Tracing BEN2 foreground",
)
import numpy as np
from PIL import Image
alpha = np.asarray(foreground.getchannel("A"), dtype=np.uint8)
contour = contour_from_alpha(alpha, selection_mask, origin)
gated_alpha = np.where(selection_mask, alpha, 0).astype(np.uint8)
foreground_pixels = foreground.convert("RGBA")
foreground_pixels.putalpha(Image.fromarray(gated_alpha, mode="L"))
foreground_pixels.save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
contour_path = output_directory / "contour.json"
contour_path.write_text(
json.dumps({"points": contour}, separators=(",", ":")),
encoding="utf-8",
)
result["contour"] = "contour.json"
if depth_analysis:
context.progress(
0.12 if not contour_analysis else 0.58,
"Estimating Balloon depth with MoGe-2",
)
prediction = moge2.infer(
model_manager,
{
"model": parameters["depth_model"],
"device": device,
"resolution_level": parameters.get("resolution_level", 5),
},
[input_path],
include_points=False,
progress_callback=None,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(0.9, "Writing Balloon depth")
depth_path = write_depth_analysis(
prediction,
selection_mask,
origin,
source_size,
output_directory / "depth.npz",
)
result["depth"] = depth_path.name
import numpy as np
with np.load(depth_path, allow_pickle=False) as data:
normal_path = write_camera_normal_texture(
prediction.normal[0],
data["valid"],
origin,
(selection_mask.shape[1], selection_mask.shape[0]),
source_size,
output_directory / "normal.png",
)
result["normal"] = normal_path.name
total_ms = (time.perf_counter() - total_started) * 1000.0
message = f"Balloon AI analysis finished in {total_ms:.0f} ms"
context.progress(1.0, message)
return result

176
src/anyimage/server/jobs/cutout_depth.py Normal file
View File

@@ -0,0 +1,176 @@
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

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

@@ -3,7 +3,7 @@ import json
from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from .background import remove_image_background
from .balloon import (
from .cutout_depth import (
write_camera_normal_texture,
write_depth_analysis,
)
@@ -247,7 +247,7 @@ def _write_camera_normal(context, parameters, color_path):
)
def _write_balloon_depth(
def _write_depth_balloon(
context,
parameters,
input_path,
@@ -371,7 +371,7 @@ def run(context, parameters):
).name
elif shape == "DEPTH_BALLOON":
context.progress(0.4, "Estimating Balloon depth")
depth_path, normal_path = _write_balloon_depth(
depth_path, normal_path = _write_depth_balloon(
context,
parameters,
model_path,