AnyImage/src/anyimage/operators/cutout_tool/main.py

"""Convert one image Region into a Cutout Shape."""

import json

import bpy

from ...preferences import (
    MAX_CUTOUT_MESH_RESOLUTION,
    MIN_CUTOUT_MESH_RESOLUTION,
    addon_preferences,
    configured_max_ai_input_size,
    configured_max_cutout_mesh_resolution,
)
from ...properties import (
    BEN2_MODEL_KEY,
    DEFAULT_MOGE2_MODEL_KEY,
    refine_region_property,
)
from ...runtime import JobOperatorBase
from ...server.media.resolution import (
    DEFAULT_MAX_AI_INPUT_SIZE,
    MAX_MAX_AI_INPUT_SIZE,
    MIN_MAX_AI_INPUT_SIZE,
)
from ..common.ai_job import (
    invoke_ai_setup_if_needed,
    moge2_parameters,
    production_device,
    require_environment,
    require_input_path,
    require_model,
)
from ..common.image_data import (
    cleanup_image_input,
    image_rgba,
    load_normal_result_image,
    prepare_region_input,
    require_image_empty,
    source_start_frame,
    video_total_frames,
)
from ..common.depth_data import load_depth_metadata, load_depth_result_image
from ..common.image_interaction import (
    ImageRegionGesture,
    apply_region_settings,
    draw_region_settings,
    gesture_property,
    report_image_edit_exception,
    resolve_image_edit_click,
    serialize_matrix,
)
from ..common.image_region import (
    ImageEditWarning,
    RegionMask,
    RegionPath,
    rasterize_region_path,
)
from .geometry import (
    CUTOUT_ALPHA_THRESHOLD,
    build_base_shape,
    limit_cutout_mesh_resolution,
)
from .interaction import open_shape_pie, validate_cutout_region
from .object import NormalMode, create_shape_object
from .shape import CUTOUT_SHAPES, DEPTH_CUTOUT_SHAPES

_HIDDEN = {"HIDDEN", "SKIP_SAVE"}


def normal_mode_for_shape(shape, generate_normal):
    if not generate_normal:
        return NormalMode.NONE
    if shape == "DEPTH_SURFACE":
        return NormalMode.OBJECT
    return NormalMode.TANGENT


def depth_mode_for_shape(shape):
    if shape == "DEPTH_SURFACE":
        return "VECTOR"
    if shape == "DEPTH_BALLOON":
        return "Z"
    return "NONE"


def create_cutout_shape(
    context,
    source_object,
    shape,
    mesh_resolution,
    region_mask,
    *,
    depth_image=None,
    depth_metadata=None,
    normal_image=None,
    normal_mode=NormalMode.NONE,
    source_rgba=None,
):
    base_shape = build_base_shape(
        context,
        source_object,
        mesh_resolution,
        region_mask,
    )
    result = create_shape_object(
        context,
        source_object,
        shape,
        mesh_resolution,
        region_mask,
        base_shape,
        depth_image=depth_image,
        depth_metadata=depth_metadata,
        normal_image=normal_image,
        normal_mode=normal_mode,
        source_rgba=source_rgba,
    )
    undo_message = f"Convert to {shape.replace('_', ' ').title()}"
    if "FINISHED" not in bpy.ops.ed.undo_push(message=undo_message):
        raise RuntimeError(f"Unable to create the {undo_message} Undo step")
    return result


def _region_mask_from_image(path, bounds):
    image = bpy.data.images.load(str(path), check_existing=False)
    try:
        return RegionMask(
            image_rgba(image)[:, :, 3],
            tuple(int(value) for value in bounds),
        )
    finally:
        bpy.data.images.remove(image, do_unlink=True)


def _generated_region_mask(result, input_path, source_bounds):
    if not isinstance(result.value, dict) or "region_image" not in result.value:
        return _region_mask_from_image(input_path, source_bounds)
    local_bounds = tuple(
        int(value)
        for value in json.loads(
            result.file("region_bounds").read_text(encoding="utf-8")
        )
    )
    bounds = (
        source_bounds[0] + local_bounds[0],
        source_bounds[1] + local_bounds[1],
        source_bounds[0] + local_bounds[2],
        source_bounds[1] + local_bounds[3],
    )
    return _region_mask_from_image(result.file("region_image"), bounds)


def cutout_geometry_settings(context):
    preferences = addon_preferences(context)
    model = getattr(preferences, "geometry_model", DEFAULT_MOGE2_MODEL_KEY)
    return (
        model,
        int(getattr(preferences, "geometry_resolution_level", 5)),
    )


class CutoutRegionToShape(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.cutout_region_to_shape"
    bl_label = "Cutout Region to Shape"
    bl_description = "Build the selected image Region as a Cutout shape"
    bl_options = {"INTERNAL"}
    job_type = "generate-cutout-artifacts"
    starting_message = "Generating Cutout artifacts"

    input_path: bpy.props.StringProperty(options=_HIDDEN)
    source_object_name: bpy.props.StringProperty(options=_HIDDEN)
    region_path_json: bpy.props.StringProperty(options=_HIDDEN)
    region_bounds_json: bpy.props.StringProperty(options=_HIDDEN)
    shape: bpy.props.StringProperty(options=_HIDDEN)
    mesh_resolution: bpy.props.IntProperty(
        default=128,
        min=MIN_CUTOUT_MESH_RESOLUTION,
        max=MAX_CUTOUT_MESH_RESOLUTION,
        options=_HIDDEN,
    )
    model: bpy.props.StringProperty(
        default=DEFAULT_MOGE2_MODEL_KEY,
        options=_HIDDEN,
    )
    resolution_level: bpy.props.IntProperty(
        default=5,
        min=0,
        max=9,
        options=_HIDDEN,
    )
    max_ai_input_size: bpy.props.IntProperty(
        default=DEFAULT_MAX_AI_INPUT_SIZE,
        min=MIN_MAX_AI_INPUT_SIZE,
        max=MAX_MAX_AI_INPUT_SIZE,
        options=_HIDDEN,
    )
    generate_normal: bpy.props.BoolProperty(default=False, options=_HIDDEN)
    refine_region: bpy.props.BoolProperty(default=False, options=_HIDDEN)

    def _needs_generation(self):
        return (
            self.refine_region
            or self.shape in DEPTH_CUTOUT_SHAPES
            or self.generate_normal
        )

    def execute(self, context):
        try:
            source_object = require_image_empty(self.source_object_name)
            if video_total_frames(source_object.data) != 1:
                raise ValueError("Cutout currently supports still images")
            if self.shape not in CUTOUT_SHAPES:
                raise ValueError(f"Unsupported Cutout shape: {self.shape}")
            region_path = RegionPath.from_json(self.region_path_json)
            source_rgba = image_rgba(source_object.data)
            region_mask = rasterize_region_path(
                tuple(source_object.data.size),
                region_path,
                source_alpha=source_rgba[:, :, 3],
                alpha_threshold=CUTOUT_ALPHA_THRESHOLD,
                antialias=True,
            )
            context.view_layer.objects.active = source_object
            source_object.select_set(True)
            crop_width = region_mask.bounds[2] - region_mask.bounds[0]
            crop_height = region_mask.bounds[3] - region_mask.bounds[1]
            self.mesh_resolution = limit_cutout_mesh_resolution(
                self.mesh_resolution,
                (crop_width, crop_height),
                configured_max_cutout_mesh_resolution(context),
            )
            self.max_ai_input_size = configured_max_ai_input_size(context)
            if not self._needs_generation():
                create_cutout_shape(
                    context,
                    source_object,
                    self.shape,
                    self.mesh_resolution,
                    region_mask,
                    source_rgba=source_rgba,
                )
                return {"FINISHED"}
            setup_result = invoke_ai_setup_if_needed()
            if setup_result is not None:
                return setup_result
            self.model, self.resolution_level = cutout_geometry_settings(context)
            self.region_bounds_json = json.dumps(
                region_mask.bounds,
                separators=(",", ":"),
            )
            self.input_path = str(prepare_region_input(source_object.data, region_mask))
        except (
            ImageEditWarning,
            OSError,
            RuntimeError,
            TypeError,
            ValueError,
        ) as error:
            report_image_edit_exception(self, error)
            return {"CANCELLED"}
        result = super().execute(context)
        if "CANCELLED" in result:
            cleanup_image_input(self.input_path, True)
        return result

    def request(self, context):
        require_environment()
        source_path = require_input_path(self.input_path)
        depth_mode = depth_mode_for_shape(self.shape)
        normal_mode = normal_mode_for_shape(self.shape, self.generate_normal)
        if self.refine_region:
            require_model(BEN2_MODEL_KEY)
        if depth_mode != "NONE" or normal_mode != NormalMode.NONE:
            require_model(self.model)
        parameters = {
            "input": str(source_path),
            "device": production_device().lower(),
            "refine_region": self.refine_region,
            "depth_mode": depth_mode,
            "normal_mode": normal_mode.value,
            "max_input_size": int(self.max_ai_input_size),
        }
        if self.refine_region:
            parameters["refine_model"] = BEN2_MODEL_KEY
        if depth_mode != "NONE" or normal_mode != NormalMode.NONE:
            start_frame = source_start_frame(self.source_object_name)
            parameters.update(
                model=self.model,
                resolution_level=self.resolution_level,
            )
            parameters.update(
                moge2_parameters(
                    context.scene.anyimage_settings,
                    source_path,
                    self.model,
                    self.resolution_level,
                    start_frame,
                )
            )
        return parameters

    def response(self, context, result):
        source_object = require_image_empty(self.source_object_name)
        source_bounds = tuple(json.loads(self.region_bounds_json))
        region_mask = _generated_region_mask(
            result,
            self.input_path,
            source_bounds,
        )
        value = result.value if isinstance(result.value, dict) else {}
        normal_image = None
        depth_image = None
        try:
            normal_mode = normal_mode_for_shape(self.shape, self.generate_normal)
            normal_key = f"{normal_mode.value.lower()}_normal"
            if normal_key in value:
                normal_image = load_normal_result_image(
                    result.file(normal_key),
                    source_object,
                )
            depth_metadata = None
            depth_mode = depth_mode_for_shape(self.shape)
            if depth_mode != "NONE":
                depth_metadata = load_depth_metadata(result.file("depth_metadata"))
                depth_image = load_depth_result_image(
                    result.file(f"{depth_mode.lower()}_depth"),
                    source_object,
                    depth_metadata,
                )
            create_cutout_shape(
                context,
                source_object,
                self.shape,
                self.mesh_resolution,
                region_mask,
                depth_image=depth_image,
                depth_metadata=depth_metadata,
                normal_image=normal_image,
                normal_mode=normal_mode,
            )
        except Exception:
            if normal_image is not None and normal_image.users == 0:
                bpy.data.images.remove(normal_image, do_unlink=True)
            if depth_image is not None and depth_image.users == 0:
                bpy.data.images.remove(depth_image, do_unlink=True)
            raise
        return f"{self.shape.replace('_', ' ').title()} is ready"

    def cleanup(self):
        cleanup_image_input(self.input_path, True)


class SelectCutoutRegion(ImageRegionGesture, bpy.types.Operator):
    bl_idname = "anyimage.select_cutout_region"
    bl_label = "Select Cutout Region"
    bl_description = "Lasso an Image Empty region, then choose its mesh shape"
    bl_options = {"INTERNAL"}
    active_tool_id = "anyimage.cutout_lasso"
    refine_region_settings = {"LASSO": "cutout_refine_region"}

    source_object_name: bpy.props.StringProperty(options={"HIDDEN", "SKIP_SAVE"})
    source_matrix_data: bpy.props.StringProperty(options={"HIDDEN", "SKIP_SAVE"})
    region_path_json: bpy.props.StringProperty(options={"HIDDEN", "SKIP_SAVE"})
    refine_region: refine_region_property()
    invert: bpy.props.BoolProperty(default=False, options={"SKIP_SAVE"})
    gesture: gesture_property()
    full_image: bpy.props.BoolProperty(default=False, options={"HIDDEN", "SKIP_SAVE"})

    preview_fill_rule = "NONZERO"
    _pie_event = None

    def _set_status(self, context):
        invert = "On" if self.invert else "Off"
        context.workspace.status_text_set(
            f"Cutout Lasso: drag and release for shapes; "
            f"F toggles Invert ({invert}); RMB or Esc cancels"
        )

    def invoke(self, context, event):
        self._pie_event = event
        if not self.full_image and event.value != "DOUBLE_CLICK":
            return super().invoke(context, event)
        source_object, should_edit = resolve_image_edit_click(
            context, event, self.report
        )
        if not should_edit:
            return {"FINISHED"} if source_object is not None else {"CANCELLED"}
        if source_object is None:
            return {"CANCELLED"}
        if video_total_frames(source_object.data) > 1:
            self.report({"WARNING"}, "Image tools only support a still image")
            return {"CANCELLED"}
        apply_region_settings(self, context)
        self.invert = False
        self.source_object_name = source_object.name
        self.source_matrix_data = serialize_matrix(source_object.matrix_world)
        width, height = (int(value) for value in source_object.data.size)
        self.region_path_json = RegionPath(
            points=((0, 0), (width, 0), (width, height), (0, height)),
            fill_rule="NONZERO",
        ).to_json()
        return self.execute(context)

    def modal(self, context, event):
        self._pie_event = event
        return super().modal(context, event)

    def execute(self, context):
        try:
            source_object = require_image_empty(self.source_object_name)
        except RuntimeError as error:
            self.report({"ERROR"}, str(error))
            return {"CANCELLED"}
        try:
            region_path = RegionPath.from_json(self.region_path_json)
            validate_cutout_region(source_object, region_path)
        except (ImageEditWarning, TypeError, ValueError) as error:
            report_image_edit_exception(self, error)
            return {"CANCELLED"}
        open_shape_pie(
            context,
            self._pie_event,
            source_object,
            region_path,
            self.refine_region,
        )
        return {"FINISHED"}


class CutoutTool(bpy.types.WorkSpaceTool):
    bl_space_type = "VIEW_3D"
    bl_context_mode = "OBJECT"
    bl_idname = "anyimage.cutout_lasso"
    bl_label = "Cutout"
    bl_description = "Lasso an Image Empty and choose a mesh shape"
    bl_icon = "ops.sculpt.lasso_mask"
    bl_operator = SelectCutoutRegion.bl_idname
    bl_keymap = (
        (
            SelectCutoutRegion.bl_idname,
            {"type": "LEFTMOUSE", "value": "DOUBLE_CLICK"},
            {"properties": [("gesture", "LASSO"), ("full_image", True)]},
        ),
        (
            SelectCutoutRegion.bl_idname,
            {"type": "LEFTMOUSE", "value": "PRESS"},
            {"properties": [("gesture", "LASSO")]},
        ),
    )

    def draw_settings(context, layout, _tool):
        settings = context.scene.anyimage_settings
        draw_region_settings(context, layout, "cutout_refine_region")
        layout.prop(settings, "cutout_point_spacing")
        layout.prop(settings, "cutout_generate_normal")