Add new node asset scripts for geometry and shader operations
- Implemented `build_mesh_plane.py` to create a geometry node group for mesh planes with customizable parameters. - Created `build_shadeless.py` for a shader node group that displays color without scene lighting while preserving shadows and reflections. - Added `validate.py` to validate the integrity of the newly created node groups and their inputs. - Introduced `validate_cutout.py` to ensure the functionality of cutout node groups and their dependencies. - Developed `validate_mesh_plane.py` to test the mesh plane creation and its modifier settings.
This commit is contained in:
255
src/anyimage/server/jobs/prepare_geometry_data.py
Normal file
255
src/anyimage/server/jobs/prepare_geometry_data.py
Normal file
@@ -0,0 +1,255 @@
|
||||
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
|
||||
Reference in New Issue
Block a user