import json import math import bpy from ...properties import ( BEN2_MODEL_KEY, ai_ready, ai_setup_label, ai_status, ) from ...properties import ( balloon_contour_analysis_property as contour_analysis_property, ) from ...properties import ( balloon_depth_analysis_property as depth_analysis_property, ) from ...properties import ( balloon_depth_model_property as depth_model_property, ) from ...properties import ( balloon_mode_property as mode_property, ) from ...properties import ( balloon_resolution_property as resolution_property, ) from ...properties import ( balloon_reuse_material_property as reuse_material_property, ) from ...properties import ( balloon_thickness_property as thickness_property, ) from ...runtime import JobOperatorBase from ..common import ( HIDDEN, cleanup_temporary_input, production_device, push_undo, require_environment, require_model, ) from ..image import ( color_data_name, is_image_empty, material_data_name, normal_data_name, plane_bounds, prepare_image_input, remove_temporary_input, video_total_frames, ) from ..material import ( IMAGE_MATERIAL_NODE_GROUP_NAME, create_material, ) from ..normal_map import bake_tangent_normal_image from ..surface import assign_vertex_uv from .depth import ( add_balloon_shape_modifier, assign_balloon_shape_attributes, balloon_depth_fields, create_balloon_base_plane, load_balloon_depth_data, ) from .geometry import INFLATION, POINT_SPACING, triangulate_balloon MIN_PATH_DISTANCE = 2.0 BALLOON_MATERIAL_PROPERTY = "anyimage_balloon_material" BALLOON_SOURCE_IMAGE_PROPERTY = "anyimage_source_image" BALLOON_NORMAL_IMAGE_PROPERTY = "anyimage_balloon_normal_image" def simplify_screen_path(path): pixels = [] for point in path: point = (float(point[0]), float(point[1])) if not pixels or math.dist(point, pixels[-1]) >= MIN_PATH_DISTANCE: pixels.append(point) return pixels def screen_path_to_image_pixels( path, region, region_data, matrix_world, local_bounds, image_size, ): from bpy_extras import view3d_utils inverse = matrix_world.inverted() x_min, x_max, y_min, y_max = local_bounds image_width, image_height = image_size local_width = x_max - x_min local_height = y_max - y_min pixels = [] for point in simplify_screen_path(path): ray_origin = inverse @ view3d_utils.region_2d_to_origin_3d( region, region_data, point, ) ray_direction = inverse.to_3x3() @ view3d_utils.region_2d_to_vector_3d( region, region_data, point, ) if abs(ray_direction.z) <= 1e-10: raise ValueError("The Image Empty is edge-on to the current view") distance = -ray_origin.z / ray_direction.z local = ray_origin + ray_direction * distance pixels.append( ( (local.x - x_min) * image_width / local_width, (y_max - local.y) * image_height / local_height, ) ) return pixels def world_plane_metric(matrix_world): from mathutils import Vector linear = matrix_world.to_3x3() basis_x = linear @ Vector((1.0, 0.0, 0.0)) basis_y = linear @ Vector((0.0, 1.0, 0.0)) basis_z = linear @ Vector((0.0, 0.0, 1.0)) metric = ( (basis_x.dot(basis_x), basis_x.dot(basis_y)), (basis_x.dot(basis_y), basis_y.dot(basis_y)), ) return metric, basis_z.length def serialize_matrix(matrix): return json.dumps( [float(matrix[row][column]) for row in range(4) for column in range(4)], separators=(",", ":"), ) def deserialize_matrix(data): from mathutils import Matrix values = json.loads(data) if len(values) != 16: raise ValueError("The stored Image Empty transform is invalid") return Matrix(tuple(tuple(values[row * 4 : row * 4 + 4]) for row in range(4))) def balloon_geometry_parameters(thickness, resolution): if thickness <= 0.0: raise ValueError("Balloon thickness must be greater than zero") if resolution <= 0.0: raise ValueError("Balloon resolution must be greater than zero") return resolution, INFLATION * thickness def data_name_matches(name, base_name): if name == base_name: return True suffix = name.removeprefix(f"{base_name}.") return suffix != name and suffix.isdigit() def marked_balloon_material(source_image, normal_required=False): return next( ( material for material in bpy.data.materials if callable(getattr(material, "get", None)) and bool(material.get(BALLOON_MATERIAL_PROPERTY)) and material.get(BALLOON_SOURCE_IMAGE_PROPERTY) is source_image and (material.get(BALLOON_NORMAL_IMAGE_PROPERTY) is not None) == normal_required ), None, ) def mark_balloon_material(material, source_image, normal_image=None): material[BALLOON_MATERIAL_PROPERTY] = True material[BALLOON_SOURCE_IMAGE_PROPERTY] = source_image if normal_image is not None: material[BALLOON_NORMAL_IMAGE_PROPERTY] = normal_image def reusable_balloon_material(source_image, normal_required=False): return marked_balloon_material(source_image, normal_required) def load_balloon_normal_image(path, source_object): image = bpy.data.images.load( str(path), check_existing=False, ) image.name = normal_data_name(source_object) image.colorspace_settings.name = "Non-Color" return image def merge_balloon_normal_image(target, update): import numpy as np if tuple(target.size) != tuple(update.size): raise ValueError("Reusable Balloon normal image dimensions do not match") target_pixels = np.empty(len(target.pixels), dtype=np.float32) update_pixels = np.empty(len(update.pixels), dtype=np.float32) target.pixels.foreach_get(target_pixels) update.pixels.foreach_get(update_pixels) target_rgba = target_pixels.reshape((-1, 4)) update_rgba = update_pixels.reshape((-1, 4)) neutral = np.asarray((128.0 / 255.0, 128.0 / 255.0, 1.0)) changed = np.max(np.abs(update_rgba[:, :3] - neutral), axis=1) > (1.5 / 255.0) target_rgba[changed] = update_rgba[changed] target.pixels.foreach_set(target_pixels) target.pack() target.update() def reusable_color_image(color_name): return next( ( image for image in bpy.data.images if data_name_matches(image.name, color_name) ), None, ) BALLOON_AXIS_TO_SOURCE = ( (0.0, -1.0, 0.0, 0.0), (0.0, 0.0, 1.0, 0.0), (-1.0, 0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0), ) def orient_balloon_vertices(vertices): """Rotate legacy XY/Z Balloon coordinates into a YZ/-X frame.""" return [(-z, -x, y) for x, y, z in vertices] def orient_balloon_world_matrix(source_matrix): from mathutils import Matrix return source_matrix @ Matrix(BALLOON_AXIS_TO_SOURCE) def orient_balloon_base_plane(local_plane): slope_x, slope_y, offset = local_plane return (slope_x, -slope_y, -offset) def extract_balloon_object( context, source_object, regions, matrix_world, algorithm, thickness=1.0, resolution=POINT_SPACING, select_result=True, reuse_material=True, depth_data=None, normal_path=None, material_image=None, bake_normal_to_tangent=False, ): image = material_image or source_object.data image_size = tuple(image.size) local_bounds = plane_bounds(source_object) world_metric, world_z_scale = world_plane_metric(matrix_world) point_spacing, inflation = balloon_geometry_parameters( thickness, resolution, ) vertices, faces, vertex_uv = triangulate_balloon( regions, image_size, local_bounds, spacing=point_spacing, world_metric=world_metric, world_z_scale=world_z_scale, inflation=INFLATION, algorithm=algorithm, ) depth_values = None local_base_plane = (0.0, 0.0, 0.0) if depth_data is not None: depth_values, local_base_plane = balloon_depth_fields( vertices, vertex_uv, depth_data, ) local_base_plane = orient_balloon_base_plane(local_base_plane) vertices = orient_balloon_vertices(vertices) source_name = source_object.name balloon_name = source_name material_name = material_data_name(source_object) color_name = color_data_name(source_object) color_image = None uses_normal = normal_path is not None material = ( reusable_balloon_material( image, normal_required=uses_normal, ) if reuse_material else None ) normal_image = None owns_color_image = False owns_normal_image = False owns_material = False mesh = None balloon = None base_plane = None try: if material is not None and not uses_normal: mark_balloon_material(material, image) if material is not None and uses_normal: normal_image = material.get(BALLOON_NORMAL_IMAGE_PROPERTY) update_normal = load_balloon_normal_image(normal_path, source_object) try: merge_balloon_normal_image(normal_image, update_normal) finally: bpy.data.images.remove(update_normal) if material is None: color_image = material_image if color_image is None and reuse_material: color_image = reusable_color_image(color_name) if color_image is None: color_image = image.copy() color_image.name = color_name owns_color_image = True if normal_path is not None: normal_image = load_balloon_normal_image( normal_path, source_object, ) normal_image.pack() owns_normal_image = True material = create_material( material_name, color_image, 0.0, surface_group_name=IMAGE_MATERIAL_NODE_GROUP_NAME, normal_image=normal_image, scene=context.scene, ) mark_balloon_material(material, image, normal_image) owns_material = True mesh = bpy.data.meshes.new(balloon_name) mesh.from_pydata(vertices, [], faces) assign_balloon_shape_attributes( mesh, vertices, depth_values, ) mesh.update() assign_vertex_uv(mesh, vertex_uv) for face in mesh.polygons: face.use_smooth = True balloon = bpy.data.objects.new(balloon_name, mesh) context.collection.objects.link(balloon) balloon.matrix_world = orient_balloon_world_matrix(matrix_world) balloon["anyimage_source_type"] = "BALLOON" balloon["anyimage_balloon_algorithm"] = algorithm balloon["anyimage_source_object"] = source_name balloon["anyimage_inflation"] = inflation balloon["anyimage_resolution"] = point_spacing balloon.data.materials.append(material) if depth_data is not None: base_plane = create_balloon_base_plane( context, balloon, vertices, local_base_plane, ) add_balloon_shape_modifier( balloon, thickness, base_plane, ) if bake_normal_to_tangent and normal_image is not None: bake_tangent_normal_image( context, source_object, balloon, normal_image, ) if depth_data is not None: balloon["anyimage_balloon_depth_analysis"] = True balloon["anyimage_reference_depth"] = depth_data["reference_depth"] except Exception: if base_plane is not None: bpy.data.objects.remove(base_plane, do_unlink=True) if balloon is not None: bpy.data.objects.remove(balloon, do_unlink=True) if mesh is not None and mesh.users == 0: bpy.data.meshes.remove(mesh) if owns_material and material is not None and material.users == 0: bpy.data.materials.remove(material) if owns_color_image and color_image is not None and color_image.users == 0: bpy.data.images.remove(color_image) if owns_normal_image and normal_image is not None and normal_image.users == 0: bpy.data.images.remove(normal_image) raise bpy.ops.object.select_all(action="DESELECT") selected_object = balloon if select_result else source_object selected_object.select_set(True) context.view_layer.objects.active = selected_object return balloon def hidden_string_property(): return bpy.props.StringProperty(options={"HIDDEN", "SKIP_SAVE"}) def meters_to_world_distance(scene, meters): unit_settings = getattr(scene, "unit_settings", None) scale_length = getattr(unit_settings, "scale_length", 1.0) if scale_length <= 0.0: scale_length = 1.0 return meters / scale_length def draw_line_strip(vertices, color, width=2.0): import gpu from gpu_extras.batch import batch_for_shader shader = gpu.shader.from_builtin("UNIFORM_COLOR") batch = batch_for_shader(shader, "LINE_STRIP", {"pos": vertices}) gpu.state.blend_set("ALPHA") gpu.state.line_width_set(width) shader.bind() shader.uniform_float("color", color) batch.draw(shader) gpu.state.line_width_set(1.0) gpu.state.blend_set("NONE") def drawing_in_region(area_pointer, region_pointer): area = getattr(bpy.context, "area", None) region = getattr(bpy.context, "region", None) return ( area is not None and region is not None and area.as_pointer() == area_pointer and region.as_pointer() == region_pointer ) class ExtractBalloon(bpy.types.Operator): bl_idname = "anyimage.extract_balloon" bl_label = "Extract Balloon" bl_description = "Build one closed Balloon mesh from a saved lasso region" bl_options = {"REGISTER", "UNDO"} mode: mode_property() thickness: thickness_property() resolution: resolution_property() source_object_name: hidden_string_property() reuse_material: reuse_material_property() contour_analysis: contour_analysis_property() depth_analysis: depth_analysis_property() depth_model: depth_model_property() source_matrix_data: hidden_string_property() selection_data: hidden_string_property() _path = None _handle = None _area_pointer = 0 _region_pointer = 0 def draw(self, context): layout = self.layout layout.prop(self, "mode") layout.prop(self, "thickness") layout.prop(self, "resolution", text="Resolution (m)") layout.prop(self, "reuse_material") status = ai_status() ready = status["ready"] contour_row = layout.row() contour_row.enabled = ready contour_row.prop(self, "contour_analysis") depth_row = layout.row() depth_row.enabled = ready depth_row.prop(self, "depth_analysis") if self.depth_analysis: model_row = layout.row() model_row.enabled = ready model_row.prop(self, "depth_model") if not ready: layout.operator( "anyimage.setup_ai_environment", text=ai_setup_label(status), icon="IMPORT", ) @classmethod def poll(cls, context): return ( context.area is not None and context.area.type == "VIEW_3D" and is_image_empty(context.active_object) ) def _draw_lasso(self): if not self._path or not drawing_in_region( self._area_pointer, self._region_pointer ): return vertices = list(self._path) if len(vertices) > 2: vertices.append(vertices[0]) draw_line_strip(vertices, (0.15, 0.65, 1.0, 0.95), 2.0) def _finish(self, context): if self._handle is not None: bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW") self._handle = None context.workspace.status_text_set(None) if context.area is not None: context.area.tag_redraw() def cancel(self, context): self._finish(context) def invoke(self, context, event): source_object = context.active_object settings = context.scene.anyimage_settings self.mode = settings.balloon_mode self.thickness = settings.balloon_thickness self.resolution = settings.balloon_resolution self.reuse_material = settings.balloon_reuse_material ready = ai_ready() self.contour_analysis = ready and settings.balloon_contour_analysis self.depth_analysis = ready and settings.balloon_depth_analysis self.depth_model = settings.balloon_depth_model self.source_object_name = source_object.name self.source_matrix_data = serialize_matrix(source_object.matrix_world) self.selection_data = "" self._path = [(event.mouse_region_x, event.mouse_region_y)] self._area_pointer = context.area.as_pointer() self._region_pointer = context.region.as_pointer() self._handle = bpy.types.SpaceView3D.draw_handler_add( self._draw_lasso, (), "WINDOW", "POST_PIXEL", ) context.workspace.status_text_set( "Balloon Lasso: release LMB to extract, RMB or Esc to cancel" ) context.window_manager.modal_handler_add(self) context.area.tag_redraw() return {"RUNNING_MODAL"} def modal(self, context, event): if event.type in {"ESC", "RIGHTMOUSE"} and event.value == "PRESS": self._finish(context) return {"CANCELLED"} if event.type == "MOUSEMOVE": self._path.append((event.mouse_region_x, event.mouse_region_y)) context.area.tag_redraw() return {"RUNNING_MODAL"} if event.type == "LEFTMOUSE" and event.value == "RELEASE": path = simplify_screen_path(self._path) if len(path) < 3: self.report({"WARNING"}, "Draw a larger lasso") self._finish(context) return {"CANCELLED"} source_object = bpy.data.objects.get(self.source_object_name) if not is_image_empty(source_object): self.report({"ERROR"}, "The source Image Empty is no longer available") self._finish(context) return {"CANCELLED"} try: matrix_world = deserialize_matrix(self.source_matrix_data) polygon = screen_path_to_image_pixels( path, context.region, context.region_data, matrix_world, plane_bounds(source_object), tuple(source_object.data.size), ) self.selection_data = json.dumps(polygon, separators=(",", ":")) except (RuntimeError, TypeError, ValueError) as error: self.report({"ERROR"}, str(error)) self._finish(context) return {"CANCELLED"} self._finish(context) return self.execute(context) return {"PASS_THROUGH"} def execute(self, context): source_object = bpy.data.objects.get(self.source_object_name) if not is_image_empty(source_object): self.report({"ERROR"}, "The source Image Empty is no longer available") return {"CANCELLED"} try: polygon = json.loads(self.selection_data) if len(polygon) < 3: raise ValueError("The stored lasso region is invalid") if self.contour_analysis or self.depth_analysis: if video_total_frames(source_object.data) > 1: raise RuntimeError("AI Analysis only supports a still image") source_path, delete_input = prepare_image_input( source_object.data, context.scene, ) result = bpy.ops.anyimage.analyze_balloon( input_path=str(source_path), delete_input=delete_input, source_object_name=self.source_object_name, source_matrix_data=self.source_matrix_data, selection_data=self.selection_data, mode=self.mode, thickness=self.thickness, resolution=self.resolution, reuse_material=self.reuse_material, contour_analysis=self.contour_analysis, depth_analysis=self.depth_analysis, depth_model=self.depth_model, ) if "CANCELLED" in result: if delete_input: remove_temporary_input(source_path) return {"CANCELLED"} return {"FINISHED"} matrix_world = deserialize_matrix(self.source_matrix_data) resolution_world = meters_to_world_distance( context.scene, self.resolution, ) balloon = extract_balloon_object( context, source_object, ((polygon,),), matrix_world, self.mode, thickness=self.thickness, resolution=resolution_world, select_result=False, reuse_material=self.reuse_material, ) balloon["anyimage_resolution_m"] = self.resolution except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} return {"FINISHED"} class AnalyzeBalloon(JobOperatorBase, bpy.types.Operator): bl_idname = "anyimage.analyze_balloon" bl_label = "Analyze Balloon" job_type = "analyze-balloon" starting_message = "Starting Balloon AI analysis" input_path: bpy.props.StringProperty(options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) source_object_name: hidden_string_property() source_matrix_data: hidden_string_property() selection_data: hidden_string_property() mode: mode_property() thickness: thickness_property() resolution: resolution_property() reuse_material: reuse_material_property() contour_analysis: bpy.props.BoolProperty(default=False, options=HIDDEN) depth_analysis: bpy.props.BoolProperty(default=False, options=HIDDEN) depth_model: depth_model_property() resolution_level: bpy.props.IntProperty( default=5, min=0, max=9, options=HIDDEN, ) def request(self, _context): require_environment() if self.contour_analysis: require_model(BEN2_MODEL_KEY) if self.depth_analysis: require_model(self.depth_model) if not self.contour_analysis and not self.depth_analysis: raise ValueError("Select Contour or Depth") polygon = json.loads(self.selection_data) if len(polygon) < 3: raise ValueError("The stored lasso region is invalid") return { "input": str(self.input_path), "contour_analysis": self.contour_analysis, "depth_analysis": self.depth_analysis, "contour_model": BEN2_MODEL_KEY, "depth_model": self.depth_model, "device": production_device().lower(), "resolution_level": self.resolution_level, "polygon": polygon, } def response(self, context, result): source_object = bpy.data.objects.get(self.source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source Image Empty is no longer available") polygon = json.loads(self.selection_data) if self.contour_analysis: payload = json.loads(result.file("contour").read_text(encoding="utf-8")) polygon = payload.get("points", ()) if len(polygon) < 3: raise RuntimeError("BEN2 did not return a valid Balloon contour") depth_data = ( load_balloon_depth_data(result.file("depth")) if self.depth_analysis else None ) matrix_world = deserialize_matrix(self.source_matrix_data) resolution_world = meters_to_world_distance( context.scene, self.resolution, ) balloon = extract_balloon_object( context, source_object, ((polygon,),), matrix_world, self.mode, thickness=self.thickness, resolution=resolution_world, select_result=False, reuse_material=self.reuse_material, depth_data=depth_data, normal_path=(result.file("normal") if self.depth_analysis else None), bake_normal_to_tangent=self.depth_analysis, ) balloon["anyimage_resolution_m"] = self.resolution if self.contour_analysis: balloon["anyimage_balloon_contour_analysis"] = True if self.depth_analysis: balloon["anyimage_balloon_depth_model"] = self.depth_model push_undo("Extract Balloon") return "AI-analyzed Balloon is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) class ActivateBalloonTool(bpy.types.Operator): bl_idname = "anyimage.activate_balloon_tool" bl_label = "Balloon Lasso Tool" bl_description = "Activate the Balloon Lasso tool in the 3D Viewport" bl_options = {"INTERNAL"} @classmethod def poll(cls, context): return ( context.area is not None and context.area.type == "VIEW_3D" and is_image_empty(context.active_object) ) def execute(self, context): return bpy.ops.wm.tool_set_by_id( name=BalloonTool.bl_idname, space_type="VIEW_3D", ) class BalloonTool(bpy.types.WorkSpaceTool): bl_space_type = "VIEW_3D" bl_context_mode = "OBJECT" bl_idname = "anyimage.balloon_lasso" bl_label = "Balloon Lasso" bl_description = ( "Draw an Image Empty region and extract one independently undoable " "Balloon per stroke" ) bl_icon = "ops.mesh.polybuild_hover" bl_cursor = "PAINT_BRUSH" bl_operator = ExtractBalloon.bl_idname bl_keymap = ( ( ExtractBalloon.bl_idname, {"type": "LEFTMOUSE", "value": "PRESS"}, None, ), ( "wm.radial_control", {"type": "S", "value": "PRESS"}, { "properties": [ ( "data_path_primary", "scene.anyimage_settings.balloon_thickness", ), ( "zoom_path", "scene.anyimage_settings.balloon_thickness_radial_zoom", ), ] }, ), ( "wm.radial_control", {"type": "D", "value": "PRESS"}, { "properties": [ ( "data_path_primary", "scene.anyimage_settings.balloon_resolution", ), ( "zoom_path", "scene.anyimage_settings.balloon_resolution_radial_zoom", ), ] }, ), ) def draw_settings(context, layout, tool): settings = context.scene.anyimage_settings layout.prop(settings, "balloon_mode", text="") layout.prop(settings, "balloon_thickness") layout.prop(settings, "balloon_resolution", text="Resolution (m)") layout.prop(settings, "balloon_reuse_material") status = ai_status() ready = status["ready"] contour_row = layout.row() contour_row.enabled = ready contour_row.prop(settings, "balloon_contour_analysis") depth_row = layout.row() depth_row.enabled = ready depth_row.prop(settings, "balloon_depth_analysis") if settings.balloon_depth_analysis: model_row = layout.row() model_row.enabled = ready model_row.prop(settings, "balloon_depth_model", text="Depth Model") if not ready: layout.operator( "anyimage.setup_ai_environment", text=ai_setup_label(status), icon="IMPORT", )