diff --git a/README.md b/README.md index e18b5ed..83a34b8 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Environment Bootstrap 使用连续的 0–100%：uv 为 1–20%、Python 为 20 时显示 `Download Required Models`，一次补齐缺失项，不会重复下载已经就绪的模型。 Environment 未安装时，同一位置显示两个安装按钮。Environment 或必需模型未就绪时， 菜单仍显示 `Convert to Depth Plane`、`Convert to Surface` 和 `Remove Background`，但三个入口都会保持禁用。 -`Balloon Mesh (Lasso)` 是 Blender 本地几何功能，不依赖 Environment 或模型，在上述状态下仍可使用。 +`Balloon Mesh (Lasso)` 默认是 Blender 本地几何功能，在上述状态下仍可使用；只有开启 `AI Refine` 时才需要 Environment 与 BEN2。 这些 UI 状态直接检查 Storage Root 的 `models/` 文件，不会仅因打开菜单或面板而启动 Job Server。 ## 快速上手 @@ -56,11 +56,11 @@ Environment 未安装时，同一位置显示两个安装按钮。Environment 按住左键沿目标轮廓绘制，松开后生成一个带原图 UV 的封闭网格。每次拖动都是一个独立完成的顶层 `Extract Balloon` Operator，工具本身继续保持激活，下一次拖动会创建另一个 Balloon；因此 `Ctrl+Z` 只撤销最近一次生成，`Ctrl+Shift+Z` 重做。每次生成后源 Image Empty 都保持选中和 Active。未绘制时没有持续 Modal 拦截事件，视图导航和其他原生快捷键由 Blender 正常处理；绘制中可用右键或 `Esc` 取消当前轨迹。选择左侧工具栏中的其他工具即可离开 Balloon Lasso。 -- `Mode`、`Thickness`、`Resolution (m)` 和 `Reuse Material` 显示在 Tool Header / Tool Settings +- `Mode`、`Thickness`、`Resolution (m)`、`Reuse Material` 和 `AI Refine` 显示在 Tool Header / Tool Settings - `S`：调用 Blender 原生 Radial Control 调节 Thickness；通过 `zoom_path` 把默认 `1.0` 映射为约 100 px 的交互圆，软上限为 `5.0`，不设硬上限 - `D`：调用 Blender 原生 Radial Control 调节 Resolution；通过独立 `zoom_path` 把默认 `0.1 m` 映射为约 100 px 的交互圆，数值越小网格越密 -`Resolution (m)` 会根据 Scene Unit Scale 换算为网格使用的世界距离，因此不同单位设置下仍代表相同的物理采样间距。两种模式共用轮廓平滑、`.color` 图片副本和 `O Texture Layer` 材质。开启 `Reuse Material` 后，连续创建会按图片基名复用同一张 `.color` 与兼容的 Balloon 材质；关闭后每次新建。Balloon、Surface 与 Depth Plane 的材质名均取自图片数据名，而不是 Image Empty 对象名。该流程不运行 BEN2 或任何其他 AI，也不需要安装 Environment。当前只支持静态图片，不处理 Movie 或 Image Sequence。 +`Resolution (m)` 会根据 Scene Unit Scale 换算为网格使用的世界距离，因此不同单位设置下仍代表相同的物理采样间距。两种模式共用轮廓平滑、`.color` 图片副本和 `O Texture Layer` 材质。开启 `Reuse Material` 后，连续创建会按图片基名复用同一张 `.color` 与兼容的 Balloon 材质；关闭后每次新建。Balloon、Surface 与 Depth Plane 的材质名均取自图片数据名，而不是 Image Empty 对象名。`AI Refine` 默认关闭，此时 Balloon 完全在 Blender 本地生成，不需要 Environment；开启后会把套索外接范围裁切为图片，并将不规则选区外的 RGB 与 Alpha 全部清零后送入 BEN2。BEN2 Alpha 严格大于 `0.5` 的最大连通前景会成为新的套索轮廓，再生成 Balloon。当前只支持静态图片，不处理 Movie 或 Image Sequence。 ### 去除背景 选中 Image Empty 后，右键打开 `SnapForm` 子菜单并选择 `Remove Background`，完成后 Empty 会自动 diff --git a/docs/features.md b/docs/features.md index 297f929..0bc02d7 100644 --- a/docs/features.md +++ b/docs/features.md @@ -20,7 +20,7 @@ Depth Plane 与 Surface 可使用 DA3 或 MoGe-2。DA3 提供相对深度和相 转换任务完成后，SnapForm 在一个 Undo 步骤中创建网格、图片、材质和修改器。源 Image Empty 是否保留由转换选项决定。 -Balloon 是 Blender 内部工具，无需 Environment 或模型。它支持连续绘制，每次 Stroke 都生成一个可单独 Undo/Redo 的结果。Tool Settings 提供模式、厚度、采样距离与材质复用设置。 +Balloon 默认在 Blender 内部生成，无需 Environment 或模型。它支持连续绘制，每次 Stroke 都生成一个可单独 Undo/Redo 的结果。Tool Settings 提供模式、厚度、采样距离、材质复用与 `AI Refine`。开启 `AI Refine` 后，Server 会裁切套索外接范围，把不规则选区外的 RGB 与 Alpha 清零，再用 BEN2 Alpha 严格大于 `0.5` 的最大连通前景重建轮廓；此模式需要 Environment 与 BEN2。 ## 动画 diff --git a/docs/outputs.md b/docs/outputs.md index c6a3531..cd1c1ec 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -9,6 +9,7 @@ | `generate-depth` | `results.npz`、深度与颜色序列 | Depth Plane 所需深度、相机和贴图 | | `generate-surface` | `mesh.npz`、`normal.png`、颜色层 | Surface 网格、相机与材质输入 | | `remove-background` | `color.0001.png` 起 | 带 Alpha 的颜色层 | +| `refine-balloon` | `lasso-input.png`、`foreground.png`、`contour.json` | 套索约束输入、BEN2 前景与重建轮廓 | | `debug-da3` | `prediction.npz` | DA3 API 返回 | | `debug-moge2` | `prediction*.npz` | MoGe-2 API 返回 | | `debug-ben2` | `foreground*.png` | BEN2 API 返回 | diff --git a/docs/server.md b/docs/server.md index 6e74df0..634862e 100644 --- a/docs/server.md +++ b/docs/server.md @@ -57,6 +57,8 @@ BlendJob 提供 Runtime、HTTP transport、队列、状态和取消协议。Snap Environment 使用 Python 3.12、ONNX Runtime、NumPy、Pillow、PyAV、trimesh、FastAPI 和 Uvicorn。Windows 使用 DirectML 或 CPU；其他平台根据支持情况使用 CoreML、OpenVINO 或 CPU。 +Balloon 的 `refine-balloon` Job 会在 BEN2 推理前裁切套索外接范围，并将不规则选区外的 RGBA 清零；推理后只回传 Alpha 大于 `0.5` 的最大连通轮廓，由 Blender 主线程生成网格。 + 正式功能支持： - MoGe-2 ViT-S、ViT-B、ViT-L Normal diff --git a/src/snapform/operators/__init__.py b/src/snapform/operators/__init__.py index 679963c..298e639 100644 --- a/src/snapform/operators/__init__.py +++ b/src/snapform/operators/__init__.py @@ -2,6 +2,7 @@ from .background import ProcessImageBackground, RemoveImageBackground from .balloon import ( ActivateBalloonTool, ExtractBalloon, + RefineBalloon, ) from .debug import DebugBen2, DebugDa3, DebugMoge2 from .depth_plane import ( @@ -30,6 +31,7 @@ CLASSES = ( ExtractOriginalFrames, ClearModels, ExtractBalloon, + RefineBalloon, ActivateBalloonTool, RemoveImageBackground, ConvertToDepthPlane, diff --git a/src/snapform/operators/balloon/__init__.py b/src/snapform/operators/balloon/__init__.py index b33ec8d..b60a402 100644 --- a/src/snapform/operators/balloon/__init__.py +++ b/src/snapform/operators/balloon/__init__.py @@ -2,4 +2,5 @@ from .operator import ( ActivateBalloonTool, BalloonTool, ExtractBalloon, + RefineBalloon, ) diff --git a/src/snapform/operators/balloon/operator.py b/src/snapform/operators/balloon/operator.py index 7a5ac69..5a9c009 100644 --- a/src/snapform/operators/balloon/operator.py +++ b/src/snapform/operators/balloon/operator.py @@ -5,16 +5,30 @@ import bpy from .geometry import INFLATION, POINT_SPACING, triangulate_balloon from ...properties import ( + BEN2_MODEL_KEY, + balloon_ai_refine_property as ai_refine_property, balloon_mode_property as mode_property, balloon_resolution_property as resolution_property, balloon_reuse_material_property as reuse_material_property, 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, plane_bounds, + prepare_image_input, + remove_temporary_input, + video_total_frames, ) from ..material import TEXTURE_MATERIAL_NODE_GROUP_NAME, create_material from ..surface import assign_vertex_uv @@ -309,6 +323,7 @@ class ExtractBalloon(bpy.types.Operator): resolution: resolution_property() source_object_name: hidden_string_property() reuse_material: reuse_material_property() + ai_refine: ai_refine_property() source_matrix_data: hidden_string_property() selection_data: hidden_string_property() @@ -323,6 +338,7 @@ class ExtractBalloon(bpy.types.Operator): layout.prop(self, "thickness") layout.prop(self, "resolution", text="Resolution (m)") layout.prop(self, "reuse_material") + layout.prop(self, "ai_refine") @classmethod def poll(cls, context): @@ -361,6 +377,7 @@ class ExtractBalloon(bpy.types.Operator): self.thickness = settings.balloon_thickness self.resolution = settings.balloon_resolution self.reuse_material = settings.balloon_reuse_material + self.ai_refine = settings.balloon_ai_refine self.source_object_name = source_object.name self.source_matrix_data = serialize_matrix(source_object.matrix_world) self.selection_data = "" @@ -432,6 +449,29 @@ class ExtractBalloon(bpy.types.Operator): polygon = json.loads(self.selection_data) if len(polygon) < 3: raise ValueError("The stored lasso region is invalid") + if self.ai_refine: + if video_total_frames(source_object.data) > 1: + raise RuntimeError("AI Refine only supports a still image") + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + result = bpy.ops.snapform.refine_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, + ) + 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, @@ -456,6 +496,69 @@ class ExtractBalloon(bpy.types.Operator): +class RefineBalloon(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.refine_balloon" + bl_label = "Refine Balloon" + job_type = "refine-balloon" + starting_message = "Starting Balloon AI refinement" + + 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() + + def request(self, _context): + require_environment() + require_model(BEN2_MODEL_KEY) + polygon = json.loads(self.selection_data) + if len(polygon) < 3: + raise ValueError("The stored lasso region is invalid") + return { + "input": str(self.input_path), + "model": BEN2_MODEL_KEY, + "device": production_device().lower(), + "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") + 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") + 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["snapform_resolution_m"] = self.resolution + balloon["snapform_balloon_ai_refined"] = True + push_undo("Extract Balloon") + return "AI-refined Balloon is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + class ActivateBalloonTool(bpy.types.Operator): bl_idname = "snapform.activate_balloon_tool" bl_label = "Balloon Lasso Tool" @@ -533,4 +636,5 @@ class BalloonTool(bpy.types.WorkSpaceTool): 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") \ No newline at end of file + layout.prop(settings, "balloon_reuse_material") + layout.prop(settings, "balloon_ai_refine") diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 13693c4..938be2b 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -228,6 +228,17 @@ def balloon_reuse_material_property(): ) +def balloon_ai_refine_property(): + return bpy.props.BoolProperty( + name="AI Refine", + description=( + "Use BEN2 Alpha above 0.5 to refine the lasso before building " + "the Balloon" + ), + default=False, + ) + + def balloon_radial_zoom_property(scale): return bpy.props.FloatVectorProperty( size=2, @@ -241,6 +252,7 @@ class SnapFormSettings(bpy.types.PropertyGroup): balloon_thickness: balloon_thickness_property() balloon_resolution: balloon_resolution_property() balloon_reuse_material: balloon_reuse_material_property() + balloon_ai_refine: balloon_ai_refine_property() balloon_thickness_radial_zoom: balloon_radial_zoom_property(100.0) balloon_resolution_radial_zoom: balloon_radial_zoom_property(1000.0) input_path: bpy.props.StringProperty( diff --git a/src/snapform/server/__init__.py b/src/snapform/server/__init__.py index 08a03e8..42e1219 100644 --- a/src/snapform/server/__init__.py +++ b/src/snapform/server/__init__.py @@ -5,6 +5,7 @@ from blendjob import JobServer from .jobs import generate_depth as _generate_depth from .jobs import generate_surface as _generate_surface from .jobs.background import run as _remove_background +from .jobs.balloon import run as _refine_balloon from .jobs.debug_ben2 import run as _debug_ben2 from .jobs.debug_da3 import run as _debug_da3 from .jobs.debug_moge2 import run as _debug_moge2 @@ -32,6 +33,11 @@ def download_required_models(context, parameters): return manager.download_missing(context, parameters["models"]) +@server.job("refine-balloon") +def refine_balloon(context, parameters): + return _refine_balloon(context, parameters) + + @server.job("remove-background") def remove_background(context, parameters): return _remove_background(context, parameters) diff --git a/src/snapform/server/jobs/balloon.py b/src/snapform/server/jobs/balloon.py new file mode 100644 index 0000000..4618603 --- /dev/null +++ b/src/snapform/server/jobs/balloon.py @@ -0,0 +1,235 @@ +import json +import math +import time +from pathlib import Path + +from ..models.ben2 import infer_foreground + + +ALPHA_THRESHOLD = 0.5 + + +def create_lasso_input(input_path, polygon, output_path): + """Crop a lasso and clear every RGBA channel outside its polygon.""" + import numpy as np + from PIL import Image, ImageDraw + + points = np.asarray(polygon, dtype=np.float64) + if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2: + raise ValueError("The lasso does not define a valid image region") + if not np.isfinite(points).all(): + raise ValueError("The lasso contains invalid image coordinates") + + with Image.open(input_path) as opened_image: + source = opened_image.convert("RGBA") + width, height = source.size + left = max(0, int(math.floor(float(points[:, 0].min())))) + top = max(0, int(math.floor(float(points[:, 1].min())))) + right = min(width, int(math.ceil(float(points[:, 0].max())))) + bottom = min(height, int(math.ceil(float(points[:, 1].max())))) + if right <= left or bottom <= top: + raise ValueError("The lasso does not overlap the source image") + + relative = points - np.asarray((left, top), dtype=np.float64) + cropped = source.crop((left, top, right, bottom)) + selection = Image.new("L", cropped.size, 0) + ImageDraw.Draw(selection).polygon( + [tuple(point) for point in relative], + fill=255, + ) + selection_mask = np.asarray(selection, dtype=np.uint8) > 0 + pixels = np.asarray(cropped, dtype=np.uint8).copy() + pixels[~selection_mask] = 0 + masked = Image.fromarray(pixels, mode="RGBA") + output_path = Path(output_path) + masked.save(output_path, format="PNG", compress_level=1) + return output_path, selection_mask, (left, top) + + +def largest_component(mask): + """Return the largest four-connected component of a boolean image.""" + import numpy as np + + mask = np.asarray(mask, dtype=bool) + if mask.ndim != 2: + raise ValueError("The BEN2 Alpha mask must be two-dimensional") + height, width = mask.shape + visited = np.zeros_like(mask) + largest = [] + for start_y, start_x in np.argwhere(mask): + if visited[start_y, start_x]: + continue + visited[start_y, start_x] = True + pending = [(int(start_x), int(start_y))] + component = [] + while pending: + x, y = pending.pop() + component.append((x, y)) + for next_x, next_y in ( + (x - 1, y), + (x + 1, y), + (x, y - 1), + (x, y + 1), + ): + if ( + 0 <= next_x < width + and 0 <= next_y < height + and mask[next_y, next_x] + and not visited[next_y, next_x] + ): + visited[next_y, next_x] = True + pending.append((next_x, next_y)) + if len(component) > len(largest): + largest = component + if not largest: + raise ValueError("BEN2 found no foreground inside the lasso") + result = np.zeros_like(mask) + component_x, component_y = zip(*largest) + result[component_y, component_x] = True + return result + + +def _choose_boundary_edge(start, direction, candidates): + right = (-direction[1], direction[0]) + left = (direction[1], -direction[0]) + back = (-direction[0], -direction[1]) + preference = {right: 0, direction: 1, left: 2, back: 3} + return min( + candidates, + key=lambda end: preference.get( + (end[0] - start[0], end[1] - start[1]), + 4, + ), + ) + + +def _simplify_axis_aligned_contour(points): + simplified = [] + for index, current in enumerate(points): + previous = points[index - 1] + following = points[(index + 1) % len(points)] + first = (current[0] - previous[0], current[1] - previous[1]) + second = (following[0] - current[0], following[1] - current[1]) + if first[0] * second[1] - first[1] * second[0] != 0: + simplified.append(current) + return simplified + + +def trace_component_contour(component): + """Trace the outer pixel-edge loop of one four-connected component.""" + import numpy as np + + component = np.asarray(component, dtype=bool) + height, width = component.shape + edges = {} + + def add_edge(start, end): + edges.setdefault(start, []).append(end) + + for y, x in np.argwhere(component): + x = int(x) + y = int(y) + if y == 0 or not component[y - 1, x]: + add_edge((x, y), (x + 1, y)) + if x == width - 1 or not component[y, x + 1]: + add_edge((x + 1, y), (x + 1, y + 1)) + if y == height - 1 or not component[y + 1, x]: + add_edge((x + 1, y + 1), (x, y + 1)) + if x == 0 or not component[y, x - 1]: + add_edge((x, y + 1), (x, y)) + if not edges: + raise ValueError("BEN2 foreground has no usable boundary") + + start = min(edges, key=lambda point: (point[1], point[0])) + current = start + direction = (1, 0) + contour = [] + edge_count = sum(len(values) for values in edges.values()) + for _index in range(edge_count + 1): + contour.append(current) + candidates = edges.get(current, ()) + if not candidates: + raise ValueError("BEN2 foreground boundary is open") + following = _choose_boundary_edge(current, direction, candidates) + direction = (following[0] - current[0], following[1] - current[1]) + current = following + if current == start: + break + else: + raise ValueError("BEN2 foreground boundary could not be closed") + contour = _simplify_axis_aligned_contour(contour) + if len(contour) < 3: + raise ValueError("BEN2 foreground is too small for a Balloon") + return contour + + +def contour_from_alpha(alpha, selection_mask, origin=(0, 0)): + """Build a lasso contour from BEN2 Alpha values strictly above 0.5.""" + import numpy as np + + alpha = np.asarray(alpha, dtype=np.float32) + selection_mask = np.asarray(selection_mask, dtype=bool) + if alpha.shape != selection_mask.shape: + raise ValueError("BEN2 Alpha does not match the lasso input") + foreground = (alpha / 255.0 > ALPHA_THRESHOLD) & selection_mask + component = largest_component(foreground) + contour = trace_component_contour(component) + offset_x, offset_y = origin + return [[x + offset_x, y + offset_y] for x, y in contour] + + +def run(context, parameters): + """Refine a Balloon lasso with BEN2 and return its largest contour.""" + total_started = time.perf_counter() + context.progress(0.02, "Preparing lasso input") + context.check_cancelled() + + output_directory = context.directory + output_directory.mkdir(parents=True, exist_ok=True) + input_path, selection_mask, origin = create_lasso_input( + parameters["input"], + parameters["polygon"], + output_directory / "lasso-input.png", + ) + context.check_cancelled() + + device = parameters["device"] + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(parameters["model"]) + context.progress(0.18, f"Loading BEN2 for {device.upper()}") + foreground, load_ms, inference_ms = infer_foreground( + input_path, + model_directory, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + + context.progress(0.9, "Tracing BEN2 foreground") + import numpy as np + from PIL import Image + + alpha = np.asarray(foreground.getchannel("A"), dtype=np.uint8) + contour = contour_from_alpha(alpha, selection_mask, origin) + gated_alpha = np.where(selection_mask, alpha, 0).astype(np.uint8) + foreground_pixels = foreground.convert("RGBA") + foreground_pixels.putalpha(Image.fromarray(gated_alpha, mode="L")) + foreground_pixels.save( + output_directory / "foreground.png", + format="PNG", + compress_level=1, + ) + contour_path = output_directory / "contour.json" + contour_path.write_text( + json.dumps({"points": contour}, separators=(",", ":")), + encoding="utf-8", + ) + + total_ms = (time.perf_counter() - total_started) * 1000.0 + message = ( + f"Balloon contour refined in {inference_ms:.0f} ms " + f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" + ) + context.progress(1.0, message) + return {"contour": "contour.json"} diff --git a/tests/test_addon.py b/tests/test_addon.py index 839bc3a..814518d 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -801,6 +801,29 @@ class AddonTest(unittest.TestCase): self.addon.CLASSES.index(self.balloon.ActivateBalloonTool), ) + def test_balloon_ai_refine_is_registered_and_forwarded_to_ben2(self): + import inspect + + self.assertIn(self.balloon.RefineBalloon, self.addon.CLASSES) + operator = self.balloon.ExtractBalloon + self.assertIn("ai_refine", operator.__annotations__) + settings = self.addon.properties.SnapFormSettings.__annotations__ + self.assertIn("balloon_ai_refine", settings) + self.assertFalse(settings["balloon_ai_refine"]["default"]) + + execute_source = inspect.getsource(operator.execute) + invoke_source = inspect.getsource(operator.invoke) + request_source = inspect.getsource(self.balloon.RefineBalloon.request) + response_source = inspect.getsource(self.balloon.RefineBalloon.response) + tool_source = inspect.getsource(self.balloon.BalloonTool.draw_settings) + self.assertIn("bpy.ops.snapform.refine_balloon", execute_source) + self.assertIn("prepare_image_input(", execute_source) + self.assertIn("settings.balloon_ai_refine", invoke_source) + self.assertIn('"polygon": polygon', request_source) + self.assertIn('result.file("contour")', response_source) + self.assertIn("extract_balloon_object(", response_source) + self.assertIn('push_undo("Extract Balloon")', response_source) + self.assertIn('"balloon_ai_refine"', tool_source) def test_balloon_redo_properties_use_mode_and_meter_resolution(self): import inspect diff --git a/tests/test_server.py b/tests/test_server.py index fd41aeb..aab8858 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -215,6 +215,7 @@ class ServerTest(unittest.TestCase): { "download-model", "download-required-models", + "refine-balloon", "remove-background", "extract-frames", "generate-depth", @@ -381,6 +382,92 @@ class ServerTest(unittest.TestCase): self.assertEqual(status["state"], "succeeded") + def test_balloon_lasso_input_crops_and_clears_outside_rgba(self): + import numpy as np + from PIL import Image + + balloon = sys.modules["server.jobs.balloon"] + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_path = root / "source.png" + output_path = root / "lasso.png" + pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8) + pixels[1, 1, 3] = 64 + Image.fromarray(pixels, mode="RGBA").save(source_path) + + _path, _selection, origin = balloon.create_lasso_input( + source_path, + ((1, 1), (5, 1), (1, 5)), + output_path, + ) + result = np.asarray(Image.open(output_path).convert("RGBA")) + + self.assertEqual(origin, (1, 1)) + self.assertEqual(result.shape, (4, 4, 4)) + self.assertEqual(tuple(result[0, 0]), (20, 40, 60, 64)) + self.assertEqual(tuple(result[3, 3]), (0, 0, 0, 0)) + + def test_balloon_job_sends_the_masked_crop_to_ben2(self): + import numpy as np + from PIL import Image + + balloon = sys.modules["server.jobs.balloon"] + source_path = self.storage_root / "source.png" + pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8) + Image.fromarray(pixels, mode="RGBA").save(source_path) + job_directory = self.storage_root / "job" + context = SimpleNamespace( + directory=job_directory, + progress=lambda *_args: None, + check_cancelled=lambda: None, + resource=lambda _name: SimpleNamespace( + directory=lambda _model: self.storage_root / "model" + ), + ) + received = [] + + def infer(input_path, *_args, **_kwargs): + image = Image.open(input_path).convert("RGBA") + received.append(np.asarray(image)) + return image, 0.0, 1.0 + + with patch.object(balloon, "infer_foreground", side_effect=infer): + result = balloon.run( + context, + { + "input": str(source_path), + "polygon": ((1, 1), (5, 1), (1, 5)), + "model": "BEN2_BASE", + "device": "cpu", + }, + ) + + self.assertEqual(result, {"contour": "contour.json"}) + self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0)) + self.assertTrue((job_directory / "lasso-input.png").is_file()) + self.assertTrue((job_directory / "foreground.png").is_file()) + self.assertTrue((job_directory / "contour.json").is_file()) + + def test_balloon_contour_uses_alpha_above_half_and_largest_component(self): + import numpy as np + + balloon = sys.modules["server.jobs.balloon"] + alpha = np.zeros((6, 7), dtype=np.uint8) + alpha[0, 0] = 255 + alpha[1:4, 2:5] = 128 + alpha[1, 1] = 127 + selection_mask = np.ones_like(alpha, dtype=bool) + + contour = balloon.contour_from_alpha( + alpha, + selection_mask, + origin=(10, 20), + ) + + self.assertEqual( + set(map(tuple, contour)), + {(12, 21), (15, 21), (15, 24), (12, 24)}, + ) if __name__ == "__main__": unittest.main()