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:
@@ -2,20 +2,25 @@
|
||||
|
||||
from blendjob import JobServer
|
||||
|
||||
from .jobs.remove_background import run as _remove_background
|
||||
from .jobs.debug_predictions import (
|
||||
run_ben2 as _debug_ben2,
|
||||
)
|
||||
from .jobs.debug_predictions import (
|
||||
run_da3 as _debug_da3,
|
||||
)
|
||||
from .jobs.debug_predictions import (
|
||||
run_moge2 as _debug_moge2,
|
||||
)
|
||||
from .jobs.debug_predictions import (
|
||||
run_upscale as _debug_upscale,
|
||||
)
|
||||
from .jobs.refine_image_mask import run as _refine_image_mask
|
||||
from .jobs.generate_depth_plane import run as _generate_depth_plane
|
||||
from .jobs.generate_region_data import run as _generate_region_data
|
||||
from .jobs.cutout import run as _generate_cutout_artifacts
|
||||
from .jobs.depth_plane import run as _generate_depth_plane_geometry
|
||||
from .jobs.refine_image_region import run as _refine_image_region
|
||||
from .jobs.remove_background import run as _remove_background
|
||||
from .jobs.upscale import run as _upscale_image
|
||||
from .model_manager import ModelManager
|
||||
|
||||
|
||||
server = JobServer("AnyImage Job Server")
|
||||
|
||||
|
||||
@@ -36,9 +41,9 @@ def download_required_models(context, parameters):
|
||||
return manager.download_missing(context, parameters["models"])
|
||||
|
||||
|
||||
@server.job("refine-image-mask")
|
||||
def refine_image_mask(context, parameters):
|
||||
return _refine_image_mask(context, parameters)
|
||||
@server.job("refine-image-region")
|
||||
def refine_image_region(context, parameters):
|
||||
return _refine_image_region(context, parameters)
|
||||
|
||||
|
||||
@server.job("remove-background")
|
||||
@@ -51,14 +56,14 @@ def upscale_image(context, parameters):
|
||||
return _upscale_image(context, parameters)
|
||||
|
||||
|
||||
@server.job("generate-depth-plane")
|
||||
def generate_depth_plane(context, parameters):
|
||||
return _generate_depth_plane(context, parameters)
|
||||
@server.job("generate-depth-plane-geometry")
|
||||
def generate_depth_plane_geometry(context, parameters):
|
||||
return _generate_depth_plane_geometry(context, parameters)
|
||||
|
||||
|
||||
@server.job("generate-region-data")
|
||||
def generate_region_data(context, parameters):
|
||||
return _generate_region_data(context, parameters)
|
||||
@server.job("generate-cutout-artifacts")
|
||||
def generate_cutout_artifacts(context, parameters):
|
||||
return _generate_cutout_artifacts(context, parameters)
|
||||
|
||||
|
||||
@server.job("debug-da3")
|
||||
@@ -79,4 +84,3 @@ def debug_ben2(context, parameters):
|
||||
@server.job("debug-upscale")
|
||||
def debug_upscale(context, parameters):
|
||||
return _debug_upscale(context, parameters)
|
||||
|
||||
|
||||
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)
|
||||
@@ -1,21 +1,21 @@
|
||||
"""Generate optional AI data for one Blender-rasterized Cutout Region."""
|
||||
"""Generate optional AI files for one Blender-rasterized Cutout Region."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from ..geometry.generate_depth import generate_depth
|
||||
from ..region import generate_selection
|
||||
from ..geometry.prediction_artifacts import generate_moge2_artifacts
|
||||
from ..region import refine_region
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.05, "Preparing Cutout Region")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
input_path = Path(parameters["input"])
|
||||
region_path = input_path
|
||||
region_image_path = input_path
|
||||
result = {}
|
||||
if parameters.get("refine_region"):
|
||||
context.progress(0.2, "Generating refined Region")
|
||||
region_path, bounds = generate_selection(
|
||||
context.progress(0.2, "Refining Region")
|
||||
region_image_path, bounds = refine_region(
|
||||
context,
|
||||
input_path,
|
||||
parameters["refine_model"],
|
||||
@@ -28,24 +28,21 @@ def run(context, parameters):
|
||||
encoding="utf-8",
|
||||
)
|
||||
result.update(
|
||||
region=region_path.name,
|
||||
region_image=region_image_path.name,
|
||||
region_bounds=bounds_path.name,
|
||||
)
|
||||
|
||||
output_depth = bool(parameters.get("output_depth"))
|
||||
depth_mode = str(parameters.get("depth_mode", "NONE"))
|
||||
normal_mode = str(parameters.get("normal_mode", "NONE"))
|
||||
if output_depth or normal_mode != "NONE":
|
||||
archive_path, normal_path = generate_depth(
|
||||
if depth_mode != "NONE" or normal_mode != "NONE":
|
||||
artifacts = generate_moge2_artifacts(
|
||||
context,
|
||||
parameters,
|
||||
region_path,
|
||||
output_depth=output_depth,
|
||||
region_image_path,
|
||||
depth_mode=depth_mode,
|
||||
normal_mode=normal_mode,
|
||||
)
|
||||
if archive_path is not None:
|
||||
result["image_geometry"] = archive_path.name
|
||||
if normal_path is not None:
|
||||
result["normal"] = normal_path.name
|
||||
result.update({key: path.name for key, path in artifacts.items()})
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Cutout Region data is ready")
|
||||
context.progress(1.0, "Cutout artifacts are ready")
|
||||
return result
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Generate depth data and preserve color for one complete image."""
|
||||
"""Generate model-space Z Depth and preserve color for one complete image."""
|
||||
|
||||
from ..geometry.generate_depth import generate_depth
|
||||
from ..geometry.prediction_artifacts import generate_moge2_artifacts
|
||||
from ..media import copy_color_images, input_frames
|
||||
|
||||
|
||||
@@ -11,12 +11,12 @@ def run(context, parameters):
|
||||
if len(images) != 1:
|
||||
raise ValueError("Depth Plane currently supports one image")
|
||||
copy_color_images(images, context.directory)
|
||||
archive_path, _normal_path = generate_depth(
|
||||
artifacts = generate_moge2_artifacts(
|
||||
context,
|
||||
parameters,
|
||||
images[0],
|
||||
output_depth=True,
|
||||
depth_mode="Z",
|
||||
)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Depth Plane data is ready")
|
||||
return {"image_geometry": archive_path.name}
|
||||
return {key: path.name for key, path in artifacts.items()}
|
||||
@@ -1,18 +1,18 @@
|
||||
import json
|
||||
import time
|
||||
|
||||
from ..region import generate_selection
|
||||
from ..region import refine_region
|
||||
|
||||
|
||||
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.progress(0.05, "Preparing image Region")
|
||||
context.check_cancelled()
|
||||
output_directory = context.directory
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
context.progress(0.2, f"Loading BEN2 for {parameters['device'].upper()}")
|
||||
output_path, bounds = generate_selection(
|
||||
output_path, bounds = refine_region(
|
||||
context,
|
||||
parameters["input"],
|
||||
parameters["model"],
|
||||
@@ -22,6 +22,6 @@ def run(context, parameters):
|
||||
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"
|
||||
message = f"Image Region refined in {elapsed_ms:.0f} ms"
|
||||
context.progress(1.0, message)
|
||||
return {"image": output_path.name, "bounds": bounds_path.name}
|
||||
return {"region_image": output_path.name, "region_bounds": bounds_path.name}
|
||||
@@ -1,14 +1,8 @@
|
||||
from . import ben2, da3, moge2, upscale
|
||||
|
||||
|
||||
def geometry_adapter(model_key):
|
||||
return moge2 if model_key.startswith("MOGE2_") else da3
|
||||
|
||||
|
||||
__all__ = (
|
||||
"ben2",
|
||||
"da3",
|
||||
"moge2",
|
||||
"upscale",
|
||||
"geometry_adapter",
|
||||
)
|
||||
|
||||
@@ -4,7 +4,7 @@ from dataclasses import dataclass
|
||||
@dataclass
|
||||
class GeometryFrame:
|
||||
depth: object
|
||||
points: object
|
||||
points: object | None
|
||||
intrinsics: object
|
||||
valid_mask: object
|
||||
confidence: object | None = None
|
||||
|
||||
@@ -92,10 +92,10 @@ def infer_one(model_manager, parameters, image_path):
|
||||
def geometry_frame(prediction):
|
||||
import numpy as np
|
||||
|
||||
if prediction.depth.shape[0] != 1 or prediction.points is None:
|
||||
raise ValueError("Image mesh generation requires one MoGe-2 point grid")
|
||||
if prediction.depth.shape[0] != 1:
|
||||
raise ValueError("Image depth generation requires one MoGe-2 frame")
|
||||
depth = prediction.depth[0]
|
||||
points = prediction.points[0]
|
||||
points = prediction.points[0] if prediction.points is not None else None
|
||||
intrinsics = prediction.intrinsics[0].copy()
|
||||
height, width = depth.shape
|
||||
if (
|
||||
@@ -109,8 +109,9 @@ def geometry_frame(prediction):
|
||||
prediction.mask[0]
|
||||
& np.isfinite(depth)
|
||||
& (depth > 0.0)
|
||||
& np.isfinite(points).all(axis=-1)
|
||||
)
|
||||
if points is not None:
|
||||
valid &= np.isfinite(points).all(axis=-1)
|
||||
return GeometryFrame(
|
||||
depth=depth,
|
||||
points=points,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
"""AI generation for rasterized image Regions."""
|
||||
|
||||
from .generate_selection import generate_selection
|
||||
from .refine_region import refine_region
|
||||
|
||||
__all__ = ("generate_selection",)
|
||||
__all__ = ("refine_region",)
|
||||
|
||||
@@ -17,14 +17,14 @@ def _crop_to_alpha(image):
|
||||
return image.crop(bounds), bounds
|
||||
|
||||
|
||||
def generate_selection(
|
||||
def refine_region(
|
||||
context,
|
||||
input_path,
|
||||
model,
|
||||
device,
|
||||
output_path,
|
||||
):
|
||||
"""Generate a refined selection from an already rasterized Region."""
|
||||
"""Refine an image Region from an already rasterized Region."""
|
||||
model_manager = context.resource("model_manager")
|
||||
foreground, _load_ms, _inference_ms = infer_foreground(
|
||||
input_path,
|
||||
Reference in New Issue
Block a user