AnyImage/src/anyimage/server/jobs/prepare_geometry_data.py

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