feat(cutout_tool): Implement cutout region to shape functionality
- Added interaction helpers for cutout region validation and shape menu interaction. - Created main operator for converting image regions into cutout shapes, including AI artifact generation. - Introduced object creation logic for cutout shapes, handling depth and normal textures. - Defined cutout shape identifiers and node groups for various shape types. - Implemented depth texture writing and prediction artifact generation for MoGe-2. - Developed server jobs for generating cutout artifacts, depth planes, and refining image regions. - Added region refinement logic to enhance gesture masks using BEN2.
This commit is contained in:
128
src/anyimage/server/geometry/depth_texture.py
Normal file
128
src/anyimage/server/geometry/depth_texture.py
Normal file
@@ -0,0 +1,128 @@
|
||||
"""Write model-space depth textures and their compact metadata."""
|
||||
|
||||
import json
|
||||
import struct
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _exr_attribute(name, kind, value):
|
||||
return (
|
||||
name.encode("ascii")
|
||||
+ b"\0"
|
||||
+ kind.encode("ascii")
|
||||
+ b"\0"
|
||||
+ struct.pack("<I", len(value))
|
||||
+ value
|
||||
)
|
||||
|
||||
|
||||
def _exr_channels():
|
||||
return b"".join(
|
||||
name.encode("ascii")
|
||||
+ b"\0"
|
||||
+ struct.pack("<iB3xii", 2, 0, 1, 1)
|
||||
for name in ("A", "B", "G", "R")
|
||||
) + b"\0"
|
||||
|
||||
|
||||
def _write_float_exr(rgba, output_path):
|
||||
"""Write one uncompressed scanline RGBA OpenEXR without extra dependencies."""
|
||||
import numpy as np
|
||||
|
||||
rgba = np.asarray(rgba, dtype=np.float32)
|
||||
if rgba.ndim != 3 or rgba.shape[2] != 4:
|
||||
raise ValueError("Depth texture must contain RGBA float pixels")
|
||||
height, width = rgba.shape[:2]
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError("Depth texture dimensions must be positive")
|
||||
box = struct.pack("<iiii", 0, 0, width - 1, height - 1)
|
||||
header = b"".join(
|
||||
(
|
||||
struct.pack("<II", 20000630, 2),
|
||||
_exr_attribute("channels", "chlist", _exr_channels()),
|
||||
_exr_attribute("compression", "compression", b"\0"),
|
||||
_exr_attribute("dataWindow", "box2i", box),
|
||||
_exr_attribute("displayWindow", "box2i", box),
|
||||
_exr_attribute("lineOrder", "lineOrder", b"\0"),
|
||||
_exr_attribute("pixelAspectRatio", "float", struct.pack("<f", 1.0)),
|
||||
_exr_attribute(
|
||||
"screenWindowCenter",
|
||||
"v2f",
|
||||
struct.pack("<ff", 0.0, 0.0),
|
||||
),
|
||||
_exr_attribute("screenWindowWidth", "float", struct.pack("<f", 1.0)),
|
||||
b"\0",
|
||||
)
|
||||
)
|
||||
row_size = width * 4 * 4
|
||||
block_size = 8 + row_size
|
||||
first_block = len(header) + height * 8
|
||||
output_path = Path(output_path)
|
||||
with output_path.open("wb") as output:
|
||||
output.write(header)
|
||||
for row in range(height):
|
||||
output.write(struct.pack("<Q", first_block + row * block_size))
|
||||
for row in range(height):
|
||||
output.write(struct.pack("<iI", row, row_size))
|
||||
for channel in (3, 2, 1, 0):
|
||||
output.write(
|
||||
np.asarray(rgba[row, :, channel], dtype="<f4").tobytes()
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
def write_vector_depth_texture(points, valid, output_path):
|
||||
"""Write camera-space XYZ and validity as a float RGBA texture."""
|
||||
import numpy as np
|
||||
|
||||
points = np.asarray(points, dtype=np.float32)
|
||||
valid = np.asarray(valid, dtype=bool)
|
||||
if points.ndim != 3 or points.shape[2] != 3 or valid.shape != points.shape[:2]:
|
||||
raise ValueError("Vector depth dimensions do not match")
|
||||
rgba = np.zeros((*valid.shape, 4), dtype=np.float32)
|
||||
rgba[valid, :3] = points[valid]
|
||||
rgba[..., 3] = valid
|
||||
return _write_float_exr(rgba, output_path)
|
||||
|
||||
|
||||
def write_z_depth_texture(depth, valid, output_path):
|
||||
"""Write camera Z and validity as a float RGBA texture."""
|
||||
import numpy as np
|
||||
|
||||
depth = np.asarray(depth, dtype=np.float32)
|
||||
valid = np.asarray(valid, dtype=bool)
|
||||
if depth.ndim != 2 or valid.shape != depth.shape:
|
||||
raise ValueError("Z depth dimensions do not match")
|
||||
rgba = np.zeros((*depth.shape, 4), dtype=np.float32)
|
||||
rgba[..., :3] = np.where(valid, depth, 0.0)[..., None]
|
||||
rgba[..., 3] = valid
|
||||
return _write_float_exr(rgba, output_path)
|
||||
|
||||
|
||||
def write_depth_metadata(frame, reference_valid, output_path):
|
||||
"""Write camera values needed to scale a model-space depth texture."""
|
||||
import numpy as np
|
||||
|
||||
depth = np.asarray(frame.depth, dtype=np.float32)
|
||||
valid = np.asarray(frame.valid_mask, dtype=bool)
|
||||
intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
|
||||
reference_valid = np.asarray(reference_valid, dtype=bool) & valid
|
||||
if depth.ndim != 2 or valid.shape != depth.shape:
|
||||
raise ValueError("Depth metadata dimensions do not match")
|
||||
if intrinsics.shape != (3, 3):
|
||||
raise ValueError("Depth metadata requires 3x3 camera intrinsics")
|
||||
if not valid.any():
|
||||
raise ValueError("MoGe-2 returned no valid depth pixels")
|
||||
if not reference_valid.any():
|
||||
raise ValueError("The image contains no valid depth pixels")
|
||||
metadata = {
|
||||
"image_size": [int(depth.shape[1]), int(depth.shape[0])],
|
||||
"intrinsics": intrinsics.tolist(),
|
||||
"reference_depth": float(np.median(depth[reference_valid])),
|
||||
}
|
||||
output_path = Path(output_path)
|
||||
output_path.write_text(
|
||||
json.dumps(metadata, separators=(",", ":")),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return output_path
|
||||
@@ -1,181 +0,0 @@
|
||||
"""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)
|
||||
@@ -43,10 +43,10 @@ def write_camera_normal_texture(
|
||||
empty_normal = (128, 128, 255)
|
||||
elif normal_mode == "OBJECT":
|
||||
converted = normal[..., (2, 0, 1)] * np.asarray(
|
||||
(-1.0, -1.0, -1.0),
|
||||
(1.0, -1.0, -1.0),
|
||||
dtype=np.float32,
|
||||
)
|
||||
empty_normal = (255, 128, 128)
|
||||
empty_normal = (0, 128, 128)
|
||||
else:
|
||||
raise ValueError(f"Unsupported normal mode: {normal_mode}")
|
||||
encoded = np.clip(
|
||||
|
||||
122
src/anyimage/server/geometry/prediction_artifacts.py
Normal file
122
src/anyimage/server/geometry/prediction_artifacts.py
Normal file
@@ -0,0 +1,122 @@
|
||||
"""Export requested depth and normal artifacts from one MoGe-2 prediction."""
|
||||
|
||||
from ..media import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
|
||||
from ..models import moge2
|
||||
from .depth_texture import (
|
||||
write_depth_metadata,
|
||||
write_vector_depth_texture,
|
||||
write_z_depth_texture,
|
||||
)
|
||||
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 _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_moge2_artifacts(
|
||||
context,
|
||||
parameters,
|
||||
input_path,
|
||||
*,
|
||||
depth_mode="NONE",
|
||||
normal_mode="NONE",
|
||||
):
|
||||
"""Run MoGe-2 once and write only the requested file artifacts."""
|
||||
depth_mode = str(depth_mode).upper()
|
||||
normal_mode = str(normal_mode).upper()
|
||||
if depth_mode not in {"NONE", "VECTOR", "Z"}:
|
||||
raise ValueError(f"Unsupported depth mode: {depth_mode}")
|
||||
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:
|
||||
prediction = moge2.infer(
|
||||
context.resource("model_manager"),
|
||||
parameters,
|
||||
[model_path],
|
||||
include_points=depth_mode == "VECTOR",
|
||||
cancel_check=context.check_cancelled,
|
||||
)
|
||||
result = {}
|
||||
if depth_mode != "NONE":
|
||||
context.progress(0.45, "Generating depth texture")
|
||||
frame = moge2.geometry_frame(prediction)
|
||||
if depth_mode == "VECTOR":
|
||||
if frame.points is None:
|
||||
raise ValueError("Vector depth requires a MoGe-2 point field")
|
||||
depth_path = write_vector_depth_texture(
|
||||
frame.points,
|
||||
frame.valid_mask,
|
||||
context.directory / "vector-depth.exr",
|
||||
)
|
||||
result["vector_depth"] = depth_path
|
||||
else:
|
||||
depth_path = write_z_depth_texture(
|
||||
frame.depth,
|
||||
frame.valid_mask,
|
||||
context.directory / "z-depth.exr",
|
||||
)
|
||||
result["z_depth"] = depth_path
|
||||
result["depth_metadata"] = write_depth_metadata(
|
||||
frame,
|
||||
_alpha_mask(input_path, frame.valid_mask.shape),
|
||||
context.directory / "depth.json",
|
||||
)
|
||||
|
||||
if normal_mode != "NONE":
|
||||
context.progress(0.72, "Generating surface normal")
|
||||
key = f"{normal_mode.lower()}_normal"
|
||||
result[key] = _write_normal(
|
||||
prediction,
|
||||
input_path,
|
||||
context.directory / f"{normal_mode.lower()}-normal.png",
|
||||
normal_mode,
|
||||
)
|
||||
return result
|
||||
finally:
|
||||
if model_path != input_path:
|
||||
model_path.unlink(missing_ok=True)
|
||||
Reference in New Issue
Block a user