feat: Implement cutout tool and depth generation features
- Added a new cutout tool for selecting regions in images, allowing users to choose mesh shapes based on selected areas. - Introduced depth generation capabilities for rasterized images, including normal data extraction. - Created jobs for generating depth planes and region data, enhancing the AI data generation workflow. - Developed a selection refinement process for cutout regions, improving accuracy in generated outputs. - Established a new region module for handling rasterized image regions and their associated AI generation tasks. - Enhanced node graph builders for asset generation, facilitating smoother integration of new features.
This commit is contained in:
181
src/anyimage/server/geometry/generate_depth.py
Normal file
181
src/anyimage/server/geometry/generate_depth.py
Normal file
@@ -0,0 +1,181 @@
|
||||
"""Generate model-derived depth and normal data for one raster image."""
|
||||
|
||||
from ..media import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
|
||||
from ..models import geometry_adapter, moge2
|
||||
from .normal_texture import write_camera_normal_texture
|
||||
|
||||
|
||||
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 _alpha_mask(image_path, shape):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(image_path) as opened:
|
||||
alpha = opened.convert("RGBA").getchannel("A")
|
||||
alpha = alpha.resize(
|
||||
(int(shape[1]), int(shape[0])),
|
||||
Image.Resampling.NEAREST,
|
||||
)
|
||||
return np.asarray(alpha, dtype=np.uint8) >= 26
|
||||
|
||||
|
||||
def transparent_fraction(image_path, alpha_threshold=26):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(image_path) as opened:
|
||||
alpha = np.asarray(
|
||||
opened.convert("RGBA").getchannel("A"),
|
||||
dtype=np.uint8,
|
||||
)
|
||||
return float(np.count_nonzero(alpha < alpha_threshold) / alpha.size)
|
||||
|
||||
|
||||
def _write_image_geometry_archive(
|
||||
frame,
|
||||
output_path,
|
||||
calibration_valid,
|
||||
region_transparent_fraction,
|
||||
):
|
||||
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")
|
||||
calibration_valid = np.asarray(calibration_valid, dtype=bool) & valid
|
||||
if not calibration_valid.any():
|
||||
raise ValueError("The depth field contains no valid Region 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),
|
||||
region_transparent_fraction=np.asarray(
|
||||
region_transparent_fraction,
|
||||
dtype=np.float32,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _write_normal(
|
||||
prediction,
|
||||
color_path,
|
||||
output_path,
|
||||
normal_mode,
|
||||
):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
normal = np.asarray(prediction.normal[0], dtype=np.float32)
|
||||
valid = _alpha_mask(color_path, normal.shape[:2])
|
||||
with Image.open(color_path) as source:
|
||||
output_size = source.size
|
||||
return write_camera_normal_texture(
|
||||
normal,
|
||||
valid,
|
||||
(0, 0),
|
||||
output_size,
|
||||
output_size,
|
||||
output_path,
|
||||
normal_mode=normal_mode,
|
||||
)
|
||||
|
||||
|
||||
def generate_depth(
|
||||
context,
|
||||
parameters,
|
||||
input_path,
|
||||
*,
|
||||
output_depth,
|
||||
normal_mode="NONE",
|
||||
):
|
||||
"""Generate requested AI data from an already prepared image Region."""
|
||||
normal_mode = str(normal_mode).upper()
|
||||
if normal_mode not in {"NONE", "TANGENT", "OBJECT"}:
|
||||
raise ValueError(f"Unsupported normal mode: {normal_mode}")
|
||||
model_path = _limited_model_input(
|
||||
input_path,
|
||||
context.directory / "model-input.png",
|
||||
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
|
||||
)
|
||||
try:
|
||||
model_manager = context.resource("model_manager")
|
||||
adapter = geometry_adapter(parameters["model"])
|
||||
prediction = None
|
||||
if output_depth:
|
||||
context.progress(0.45, "Generating depth")
|
||||
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,
|
||||
)
|
||||
frame = adapter.geometry_frame(prediction)
|
||||
calibration_valid = _alpha_mask(input_path, frame.valid_mask.shape)
|
||||
archive_path = context.directory / "image_geometry.npz"
|
||||
_write_image_geometry_archive(
|
||||
frame,
|
||||
archive_path,
|
||||
calibration_valid,
|
||||
transparent_fraction(input_path),
|
||||
)
|
||||
else:
|
||||
archive_path = None
|
||||
|
||||
normal_path = None
|
||||
if normal_mode != "NONE":
|
||||
context.progress(0.72, "Generating surface normal")
|
||||
if prediction is None or adapter is not moge2:
|
||||
prediction = moge2.infer(
|
||||
model_manager,
|
||||
parameters,
|
||||
[model_path],
|
||||
include_points=False,
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
normal_path = _write_normal(
|
||||
prediction,
|
||||
input_path,
|
||||
context.directory / "normal.png",
|
||||
normal_mode,
|
||||
)
|
||||
return archive_path, normal_path
|
||||
finally:
|
||||
if model_path != input_path:
|
||||
model_path.unlink(missing_ok=True)
|
||||
@@ -3,14 +3,14 @@
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def resized_selection_mask(selection_mask, shape):
|
||||
def resized_valid_mask(valid_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")
|
||||
valid_mask = np.asarray(valid_mask, dtype=bool)
|
||||
if valid_mask.shape == tuple(shape):
|
||||
return valid_mask
|
||||
resized = Image.fromarray(valid_mask.astype(np.uint8) * 255, mode="L")
|
||||
resized = resized.resize(
|
||||
(int(shape[1]), int(shape[0])),
|
||||
Image.Resampling.NEAREST,
|
||||
@@ -25,8 +25,9 @@ def write_camera_normal_texture(
|
||||
crop_size,
|
||||
source_size,
|
||||
output_path,
|
||||
normal_mode="TANGENT",
|
||||
):
|
||||
"""Encode MoGe camera-axis normals into a full-image RGB texture."""
|
||||
"""Encode MoGe camera-axis normals in Blender tangent or object space."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
@@ -35,14 +36,26 @@ def write_camera_normal_texture(
|
||||
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])
|
||||
valid = resized_valid_mask(valid, normal.shape[:2])
|
||||
normal_mode = str(normal_mode).upper()
|
||||
if normal_mode == "TANGENT":
|
||||
converted = normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32)
|
||||
empty_normal = (128, 128, 255)
|
||||
elif normal_mode == "OBJECT":
|
||||
converted = normal[..., (2, 0, 1)] * np.asarray(
|
||||
(-1.0, -1.0, -1.0),
|
||||
dtype=np.float32,
|
||||
)
|
||||
empty_normal = (255, 128, 128)
|
||||
else:
|
||||
raise ValueError(f"Unsupported normal mode: {normal_mode}")
|
||||
encoded = np.clip(
|
||||
(normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32) + 1.0)
|
||||
(converted + 1.0)
|
||||
* 127.5,
|
||||
0.0,
|
||||
255.0,
|
||||
).astype(np.uint8)
|
||||
encoded[~valid] = (128, 128, 255)
|
||||
encoded[~valid] = empty_normal
|
||||
crop_width, crop_height = (int(value) for value in crop_size)
|
||||
crop = Image.fromarray(encoded, mode="RGB")
|
||||
if crop.size != (crop_width, crop_height):
|
||||
@@ -51,7 +64,7 @@ def write_camera_normal_texture(
|
||||
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 = Image.new("RGB", (source_width, source_height), empty_normal)
|
||||
full.paste(crop, tuple(int(value) for value in origin))
|
||||
output_path = Path(output_path)
|
||||
full.save(output_path)
|
||||
|
||||
Reference in New Issue
Block a user