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

import bpy

from ...properties import (
    refine_selection_property,
)
from ..common.image_data import image_rgba, require_image_empty, video_total_frames
from ..common.image_interaction import (
    apply_region_settings,
    draw_region_settings,
    report_image_edit_exception,
    resolve_image_edit_click,
)
from ..common.image_interaction import serialize_matrix
from ..common.image_interaction import ImageRegionGesture, gesture_property
from ..common.image_region import (
    ImageEditWarning,
    RegionPath,
    rasterize_region_path,
)

CUTOUT_DETAIL_MIN = 4
CUTOUT_DETAIL_MAX = 1024


def _validate_cutout_region(
    source_object,
    region_path,
):
    image = source_object.data
    rgba = image_rgba(image)
    rasterize_region_path(
        tuple(image.size),
        region_path,
        source_alpha=rgba[:, :, 3],
        alpha_threshold=0.1,
        antialias=True,
    )


def cutout_mesh_detail(image_size, point_spacing):
    longest_edge = max(int(value) for value in image_size)
    detail = round(longest_edge / max(int(point_spacing), 1))
    return max(CUTOUT_DETAIL_MIN, min(CUTOUT_DETAIL_MAX, detail))


def _pie_button(
    pie,
    shape,
    label,
    icon,
    source_name,
    region_path_json,
    refine_selection,
    normal_mode,
    mesh_detail,
):
    operator = pie.operator(
        "anyimage.cutout_region_to_shape",
        text=label,
        icon=icon,
    )
    operator.source_object_name = source_name
    operator.region_path_json = region_path_json
    operator.shape = shape
    operator.refine_selection = refine_selection
    operator.normal_mode = normal_mode
    operator.mesh_detail = mesh_detail


def _open_shape_pie(
    context,
    event,
    source_object,
    region_path,
    refine_selection=False,
):
    source_name = source_object.name
    region_path_json = region_path.to_json()
    settings = context.scene.anyimage_settings
    point_spacing = settings.cutout_point_spacing
    normal_mode = str(getattr(settings, "cutout_normal_mode", "NONE"))
    mesh_detail = cutout_mesh_detail(source_object.data.size, point_spacing)

    def draw(menu, _context):
        pie = menu.layout.menu_pie()
        _pie_button(
            pie,
            "SURFACE",
            "Surface",
            "MESH_GRID",
            source_name,
            region_path_json,
            refine_selection,
            normal_mode,
            mesh_detail,
        )
        _pie_button(
            pie,
            "BALLOON",
            "Balloon",
            "MESH_UVSPHERE",
            source_name,
            region_path_json,
            refine_selection,
            normal_mode,
            mesh_detail,
        )
        _pie_button(
            pie,
            "DEPTH_BALLOON",
            "Depth Balloon",
            "META_BALL",
            source_name,
            region_path_json,
            refine_selection,
            normal_mode,
            mesh_detail,
        )
        _pie_button(
            pie,
            "DEPTH_SURFACE",
            "Depth Surface",
            "SURFACE_NCURVE",
            source_name,
            region_path_json,
            refine_selection,
            normal_mode,
            mesh_detail,
        )

    context.window_manager.popup_menu_pie(
        event,
        draw,
        title="Cutout Shape",
        icon="MESH_DATA",
    )


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"

    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_selection: refine_selection_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_selection,
        )
        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
        layout.prop(settings, "cutout_point_spacing")
        layout.prop(settings, "cutout_normal_mode")
        draw_region_settings(context, layout, "cutout_refine_selection")