feat(cutout_tool): Implement cutout region to shape functionality
- Added interaction helpers for cutout region validation and shape menu interaction. - Created main operator for converting image regions into cutout shapes, including AI artifact generation. - Introduced object creation logic for cutout shapes, handling depth and normal textures. - Defined cutout shape identifiers and node groups for various shape types. - Implemented depth texture writing and prediction artifact generation for MoGe-2. - Developed server jobs for generating cutout artifacts, depth planes, and refining image regions. - Added region refinement logic to enhance gesture masks using BEN2.
This commit is contained in:
455
src/anyimage/operators/cutout_tool/main.py
Normal file
455
src/anyimage/operators/cutout_tool/main.py
Normal file
@@ -0,0 +1,455 @@
|
||||
"""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")
|
||||
Reference in New Issue
Block a user