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:
22
src/anyimage/server/jobs/generate_depth_plane.py
Normal file
22
src/anyimage/server/jobs/generate_depth_plane.py
Normal file
@@ -0,0 +1,22 @@
|
||||
"""Generate depth data and preserve color for one complete image."""
|
||||
|
||||
from ..geometry.generate_depth import generate_depth
|
||||
from ..media import copy_color_images, input_frames
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
context.progress(0.05, "Preparing Depth Plane")
|
||||
context.directory.mkdir(parents=True, exist_ok=True)
|
||||
with input_frames(context, parameters) as images:
|
||||
if len(images) != 1:
|
||||
raise ValueError("Depth Plane currently supports one image")
|
||||
copy_color_images(images, context.directory)
|
||||
archive_path, _normal_path = generate_depth(
|
||||
context,
|
||||
parameters,
|
||||
images[0],
|
||||
output_depth=True,
|
||||
)
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Depth Plane data is ready")
|
||||
return {"image_geometry": archive_path.name}
|
||||
51
src/anyimage/server/jobs/generate_region_data.py
Normal file
51
src/anyimage/server/jobs/generate_region_data.py
Normal file
@@ -0,0 +1,51 @@
|
||||
"""Generate optional AI data for one Blender-rasterized Cutout Region."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from ..geometry.generate_depth import generate_depth
|
||||
from ..region import generate_selection
|
||||
|
||||
|
||||
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
|
||||
result = {}
|
||||
if parameters.get("refine_region"):
|
||||
context.progress(0.2, "Generating refined Region")
|
||||
region_path, bounds = generate_selection(
|
||||
context,
|
||||
input_path,
|
||||
parameters["refine_model"],
|
||||
parameters["device"],
|
||||
context.directory / "region.png",
|
||||
)
|
||||
bounds_path = context.directory / "region_bounds.json"
|
||||
bounds_path.write_text(
|
||||
json.dumps(bounds, separators=(",", ":")),
|
||||
encoding="utf-8",
|
||||
)
|
||||
result.update(
|
||||
region=region_path.name,
|
||||
region_bounds=bounds_path.name,
|
||||
)
|
||||
|
||||
output_depth = bool(parameters.get("output_depth"))
|
||||
normal_mode = str(parameters.get("normal_mode", "NONE"))
|
||||
if output_depth or normal_mode != "NONE":
|
||||
archive_path, normal_path = generate_depth(
|
||||
context,
|
||||
parameters,
|
||||
region_path,
|
||||
output_depth=output_depth,
|
||||
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
|
||||
context.check_cancelled()
|
||||
context.progress(1.0, "Cutout Region data is ready")
|
||||
return result
|
||||
@@ -1,255 +0,0 @@
|
||||
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
|
||||
@@ -1,12 +1,7 @@
|
||||
import json
|
||||
import time
|
||||
|
||||
from ..models.ben2 import infer_foreground
|
||||
from ..selection.polygon_mask import (
|
||||
apply_selection_mask,
|
||||
create_selection_canvas,
|
||||
crop_to_alpha,
|
||||
)
|
||||
from ..region import generate_selection
|
||||
|
||||
|
||||
def run(context, parameters):
|
||||
@@ -16,28 +11,14 @@ def run(context, parameters):
|
||||
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"]),
|
||||
output_path, bounds = generate_selection(
|
||||
context,
|
||||
parameters["input"],
|
||||
parameters["model"],
|
||||
parameters["device"],
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
preserve_input_alpha=True,
|
||||
output_directory / "image.png",
|
||||
)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user