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:
@@ -10,7 +10,8 @@ from .jobs.debug_predictions import (
|
||||
run_upscale as _debug_upscale,
|
||||
)
|
||||
from .jobs.refine_image_mask import run as _refine_image_mask
|
||||
from .jobs.prepare_geometry_data import run as _prepare_geometry_data
|
||||
from .jobs.generate_depth_plane import run as _generate_depth_plane
|
||||
from .jobs.generate_region_data import run as _generate_region_data
|
||||
from .jobs.upscale import run as _upscale_image
|
||||
from .model_manager import ModelManager
|
||||
|
||||
@@ -50,9 +51,14 @@ def upscale_image(context, parameters):
|
||||
return _upscale_image(context, parameters)
|
||||
|
||||
|
||||
@server.job("prepare-geometry-data")
|
||||
def prepare_geometry_data(context, parameters):
|
||||
return _prepare_geometry_data(context, parameters)
|
||||
@server.job("generate-depth-plane")
|
||||
def generate_depth_plane(context, parameters):
|
||||
return _generate_depth_plane(context, parameters)
|
||||
|
||||
|
||||
@server.job("generate-region-data")
|
||||
def generate_region_data(context, parameters):
|
||||
return _generate_region_data(context, parameters)
|
||||
|
||||
|
||||
@server.job("debug-da3")
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
5
src/anyimage/server/region/__init__.py
Normal file
5
src/anyimage/server/region/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""AI generation for rasterized image Regions."""
|
||||
|
||||
from .generate_selection import generate_selection
|
||||
|
||||
__all__ = ("generate_selection",)
|
||||
40
src/anyimage/server/region/generate_selection.py
Normal file
40
src/anyimage/server/region/generate_selection.py
Normal file
@@ -0,0 +1,40 @@
|
||||
from ..models.ben2 import infer_foreground
|
||||
|
||||
|
||||
def _crop_to_alpha(image):
|
||||
import numpy as np
|
||||
|
||||
pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8)
|
||||
visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0)
|
||||
if not len(visible_x):
|
||||
raise ValueError("The generated Region contains no visible pixels")
|
||||
bounds = (
|
||||
int(visible_x.min()),
|
||||
int(visible_y.min()),
|
||||
int(visible_x.max()) + 1,
|
||||
int(visible_y.max()) + 1,
|
||||
)
|
||||
return image.crop(bounds), bounds
|
||||
|
||||
|
||||
def generate_selection(
|
||||
context,
|
||||
input_path,
|
||||
model,
|
||||
device,
|
||||
output_path,
|
||||
):
|
||||
"""Generate a refined selection from an already rasterized Region."""
|
||||
model_manager = context.resource("model_manager")
|
||||
foreground, _load_ms, _inference_ms = infer_foreground(
|
||||
input_path,
|
||||
model_manager.directory(model),
|
||||
device,
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
context.check_cancelled()
|
||||
output, bounds = _crop_to_alpha(foreground)
|
||||
output.save(output_path, format="PNG", compress_level=1)
|
||||
return output_path, bounds
|
||||
@@ -1,22 +0,0 @@
|
||||
"""Image selection rasterization and Cutout image preparation."""
|
||||
|
||||
from .cutout_images import refine_selection, write_images, write_selection
|
||||
from .polygon_mask import (
|
||||
apply_selection_mask,
|
||||
create_selection_canvas,
|
||||
crop_to_alpha,
|
||||
polygon_bounds,
|
||||
polygon_selection_mask,
|
||||
)
|
||||
|
||||
|
||||
__all__ = (
|
||||
"apply_selection_mask",
|
||||
"create_selection_canvas",
|
||||
"crop_to_alpha",
|
||||
"polygon_bounds",
|
||||
"polygon_selection_mask",
|
||||
"refine_selection",
|
||||
"write_images",
|
||||
"write_selection",
|
||||
)
|
||||
@@ -1,162 +0,0 @@
|
||||
"""Prepare Cutout selection images inside the Job Server."""
|
||||
|
||||
import json
|
||||
|
||||
from ..models.ben2 import infer_foreground
|
||||
from .polygon_mask import (
|
||||
apply_selection_mask,
|
||||
create_selection_canvas,
|
||||
polygon_selection_mask,
|
||||
)
|
||||
|
||||
|
||||
def _write_masked_color(input_path, polygon, output_path, invert=False):
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter, ImageOps
|
||||
|
||||
with Image.open(input_path) as opened:
|
||||
source = opened.convert("RGBA")
|
||||
selection = Image.fromarray(
|
||||
polygon_selection_mask(source.size, polygon, "NONZERO").astype(np.uint8) * 255,
|
||||
mode="L",
|
||||
).filter(ImageFilter.GaussianBlur(radius=0.75))
|
||||
if invert:
|
||||
selection = ImageOps.invert(selection)
|
||||
pixels = np.asarray(source, dtype=np.uint8).copy()
|
||||
source_alpha = pixels[:, :, 3].copy()
|
||||
pixels[:, :, 3] = np.asarray(selection, dtype=np.uint8)
|
||||
pixels[source_alpha < 26] = 0
|
||||
pixels[pixels[:, :, 3] == 0, :3] = 0
|
||||
Image.fromarray(pixels, mode="RGBA").save(
|
||||
output_path,
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
def write_selection(input_path, polygon, output_path, invert=False):
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
|
||||
if invert:
|
||||
_write_masked_color(input_path, polygon, output_path, invert=True)
|
||||
with Image.open(input_path) as source:
|
||||
return output_path, (0, 0, source.width, source.height)
|
||||
canvas, selection_mask, bounds = create_selection_canvas(
|
||||
input_path, polygon, "NONZERO"
|
||||
)
|
||||
selection = Image.fromarray(
|
||||
selection_mask.astype(np.uint8) * 255,
|
||||
mode="L",
|
||||
).filter(ImageFilter.GaussianBlur(radius=0.75))
|
||||
pixels = np.asarray(canvas, dtype=np.uint8).copy()
|
||||
source_alpha = pixels[:, :, 3].copy()
|
||||
pixels[:, :, 3] = np.asarray(selection, dtype=np.uint8)
|
||||
pixels[source_alpha < 26] = 0
|
||||
pixels[pixels[:, :, 3] == 0, :3] = 0
|
||||
Image.fromarray(pixels, mode="RGBA").save(
|
||||
output_path,
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
return output_path, bounds
|
||||
|
||||
|
||||
def refine_selection(context, parameters, input_path, polygon, output_path):
|
||||
from PIL import Image, ImageChops
|
||||
|
||||
canvas, selection_mask, bounds = create_selection_canvas(
|
||||
input_path, polygon, "NONZERO"
|
||||
)
|
||||
selection_path = context.directory / "selection-input.png"
|
||||
canvas.save(selection_path, format="PNG", compress_level=1)
|
||||
model_manager = context.resource("model_manager")
|
||||
foreground, _load_ms, _inference_ms = infer_foreground(
|
||||
selection_path,
|
||||
model_manager.directory(parameters["refine_model"]),
|
||||
parameters["device"],
|
||||
model_cache=model_manager,
|
||||
cancel_check=context.check_cancelled,
|
||||
preserve_input_alpha=True,
|
||||
)
|
||||
refined = apply_selection_mask(foreground, selection_mask)
|
||||
with Image.open(input_path) as opened:
|
||||
source = opened.convert("RGBA")
|
||||
if parameters.get("invert"):
|
||||
mask = Image.new("L", source.size, 0)
|
||||
mask.paste(refined.getchannel("A"), bounds[:2])
|
||||
source.putalpha(
|
||||
ImageChops.multiply(source.getchannel("A"), ImageChops.invert(mask))
|
||||
)
|
||||
output = source
|
||||
else:
|
||||
output = Image.new("RGBA", source.size, (0, 0, 0, 0))
|
||||
output.paste(refined, bounds[:2])
|
||||
output.save(output_path, format="PNG", compress_level=1)
|
||||
return output_path
|
||||
|
||||
|
||||
def write_images(input_path, selection_path, output_directory, selection_bounds=None):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(input_path) as opened:
|
||||
source = opened.convert("RGBA")
|
||||
with Image.open(selection_path) as opened:
|
||||
selection = opened.convert("RGBA")
|
||||
if selection_bounds is None and selection.size != source.size:
|
||||
raise ValueError("The Cutout selection does not match the source image")
|
||||
if selection_bounds is None:
|
||||
selection_bounds = (0, 0, source.width, source.height)
|
||||
color_canvas = Image.new("RGBA", selection.size, (0, 0, 0, 0))
|
||||
color_canvas.paste(source, (-selection_bounds[0], -selection_bounds[1]))
|
||||
source_alpha = np.asarray(color_canvas, dtype=np.uint8)[:, :, 3]
|
||||
selection_pixels = np.asarray(selection, dtype=np.uint8).copy()
|
||||
alpha = selection_pixels[:, :, 3]
|
||||
alpha[alpha <= 1] = 0
|
||||
alpha[source_alpha < 26] = 0
|
||||
selection_pixels[alpha == 0, :3] = 0
|
||||
selection = Image.fromarray(selection_pixels, mode="RGBA")
|
||||
visible_y, visible_x = np.nonzero(alpha >= 26)
|
||||
if not len(visible_x):
|
||||
raise ValueError("The Cutout selection contains no visible pixels")
|
||||
content_bounds = (
|
||||
int(visible_x.min()),
|
||||
int(visible_y.min()),
|
||||
int(visible_x.max()) + 1,
|
||||
int(visible_y.max()) + 1,
|
||||
)
|
||||
selection_left, selection_top, _selection_right, _selection_bottom = (
|
||||
selection_bounds
|
||||
)
|
||||
bounds = (
|
||||
selection_left + content_bounds[0],
|
||||
selection_top + content_bounds[1],
|
||||
selection_left + content_bounds[2],
|
||||
selection_top + content_bounds[3],
|
||||
)
|
||||
color = color_canvas.crop(content_bounds)
|
||||
color_path = output_directory / "image.0001.png"
|
||||
color.save(color_path, format="PNG", compress_level=1)
|
||||
|
||||
selection_crop = selection.crop(content_bounds)
|
||||
selection_path = output_directory / "selection.png"
|
||||
selection_crop.save(selection_path, format="PNG", compress_level=1)
|
||||
|
||||
analysis_pixels = np.asarray(color, dtype=np.uint8).copy()
|
||||
mask_alpha = np.asarray(selection_crop, dtype=np.uint8)[:, :, 3]
|
||||
analysis_pixels[:, :, 3] = np.minimum(analysis_pixels[:, :, 3], mask_alpha)
|
||||
analysis_pixels[mask_alpha < 26] = 0
|
||||
analysis_path = output_directory / "analysis.png"
|
||||
Image.fromarray(analysis_pixels, mode="RGBA").save(
|
||||
analysis_path,
|
||||
format="PNG",
|
||||
compress_level=1,
|
||||
)
|
||||
bounds_path = output_directory / "bounds.json"
|
||||
bounds_path.write_text(
|
||||
json.dumps(bounds, separators=(",", ":")),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return color_path, selection_path, analysis_path, bounds_path
|
||||
@@ -1,151 +0,0 @@
|
||||
import math
|
||||
|
||||
|
||||
def polygon_bounds(polygon):
|
||||
import numpy as np
|
||||
|
||||
points = np.asarray(polygon, dtype=np.float64)
|
||||
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
|
||||
raise ValueError("The selection does not define a valid image region")
|
||||
if not np.isfinite(points).all():
|
||||
raise ValueError("The selection contains invalid image coordinates")
|
||||
left = int(math.floor(float(points[:, 0].min())))
|
||||
top = int(math.floor(float(points[:, 1].min())))
|
||||
right = int(math.ceil(float(points[:, 0].max())))
|
||||
bottom = int(math.ceil(float(points[:, 1].max())))
|
||||
if right <= left or bottom <= top:
|
||||
raise ValueError("The selection does not define a valid image region")
|
||||
return left, top, right, bottom
|
||||
|
||||
|
||||
def polygon_selection_mask(image_size, polygon, fill_rule="EVEN_ODD"):
|
||||
"""Rasterize a polygon with an explicit fill rule."""
|
||||
import numpy as np
|
||||
|
||||
width, height = (int(value) for value in image_size)
|
||||
points = np.asarray(polygon, dtype=np.float64)
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError("The selection canvas has no readable dimensions")
|
||||
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
|
||||
raise ValueError("The selection does not define a valid image region")
|
||||
if not np.isfinite(points).all():
|
||||
raise ValueError("The selection contains invalid image coordinates")
|
||||
if fill_rule not in {"EVEN_ODD", "NONZERO"}:
|
||||
raise ValueError(f"Unknown polygon fill rule: {fill_rule}")
|
||||
|
||||
scanlines = np.arange(height, dtype=np.float64) + 0.5
|
||||
crossings = [[] for _index in range(height)]
|
||||
for start, end in zip(points, np.roll(points, -1, axis=0)):
|
||||
selected = np.flatnonzero(
|
||||
(start[1] > scanlines) != (end[1] > scanlines)
|
||||
)
|
||||
if not len(selected):
|
||||
continue
|
||||
crossing_x = start[0] + (
|
||||
(scanlines[selected] - start[1])
|
||||
* (end[0] - start[0])
|
||||
/ (end[1] - start[1])
|
||||
)
|
||||
winding = 1 if end[1] > start[1] else -1
|
||||
for row, value in zip(selected, crossing_x):
|
||||
crossings[int(row)].append((float(value), winding))
|
||||
|
||||
mask = np.zeros((height, width), dtype=bool)
|
||||
for row, values in enumerate(crossings):
|
||||
values.sort()
|
||||
state = 0
|
||||
index = 0
|
||||
while index < len(values):
|
||||
crossing_x = values[index][0]
|
||||
winding = 0
|
||||
count = 0
|
||||
while index < len(values) and values[index][0] == crossing_x:
|
||||
winding += values[index][1]
|
||||
count += 1
|
||||
index += 1
|
||||
state = (
|
||||
state + winding
|
||||
if fill_rule == "NONZERO"
|
||||
else state ^ (count % 2)
|
||||
)
|
||||
if index >= len(values) or state == 0:
|
||||
continue
|
||||
following_x = values[index][0]
|
||||
first = max(0, int(math.ceil(crossing_x - 0.5)))
|
||||
last = min(width, int(math.floor(following_x - 0.5)) + 1)
|
||||
mask[row, first:last] = True
|
||||
return mask
|
||||
|
||||
def create_selection_canvas(input_path, polygon, fill_rule="EVEN_ODD"):
|
||||
"""Create an unbounded selection canvas with transparent black outside."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(input_path) as opened_image:
|
||||
source = opened_image.convert("RGBA")
|
||||
left, top, right, bottom = polygon_bounds(polygon)
|
||||
if right <= 0 or bottom <= 0 or left >= source.width or top >= source.height:
|
||||
raise ValueError("The selection does not overlap the source image")
|
||||
canvas = Image.new("RGBA", (right - left, bottom - top), (0, 0, 0, 0))
|
||||
canvas.paste(source, (-left, -top))
|
||||
relative = [(float(x) - left, float(y) - top) for x, y in polygon]
|
||||
selection_mask = polygon_selection_mask(canvas.size, relative, fill_rule)
|
||||
pixels = np.asarray(canvas, dtype=np.uint8).copy()
|
||||
pixels[~selection_mask] = 0
|
||||
return Image.fromarray(pixels, mode="RGBA"), selection_mask, (
|
||||
left,
|
||||
top,
|
||||
right,
|
||||
bottom,
|
||||
)
|
||||
|
||||
|
||||
def apply_selection_mask(foreground, selection_mask):
|
||||
"""Clear pixels outside the source selection after background removal."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
pixels = np.asarray(foreground.convert("RGBA"), dtype=np.uint8).copy()
|
||||
if pixels.shape[:2] != selection_mask.shape:
|
||||
pixels = np.asarray(
|
||||
foreground.convert("RGBA").resize(
|
||||
(selection_mask.shape[1], selection_mask.shape[0]),
|
||||
Image.Resampling.BILINEAR,
|
||||
),
|
||||
dtype=np.uint8,
|
||||
)
|
||||
pixels[~selection_mask] = 0
|
||||
pixels[pixels[:, :, 3] == 0, :3] = 0
|
||||
return Image.fromarray(pixels, mode="RGBA")
|
||||
|
||||
|
||||
def crop_to_alpha(image, bounds):
|
||||
"""Crop an RGBA mask image to its nonzero Alpha bounds."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8)
|
||||
visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0)
|
||||
if not len(visible_x):
|
||||
raise ValueError("The selection contains no visible image pixels")
|
||||
content_left = int(visible_x.min())
|
||||
content_top = int(visible_y.min())
|
||||
content_right = int(visible_x.max()) + 1
|
||||
content_bottom = int(visible_y.max()) + 1
|
||||
output = Image.new(
|
||||
"RGBA",
|
||||
(
|
||||
content_right - content_left,
|
||||
content_bottom - content_top,
|
||||
),
|
||||
(0, 0, 0, 0),
|
||||
)
|
||||
output.paste(image, (-content_left, -content_top))
|
||||
left, top, _right, _bottom = bounds
|
||||
return output, (
|
||||
left + content_left,
|
||||
top + content_top,
|
||||
left + content_right,
|
||||
top + content_bottom,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user