"""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")