feat: Add AI Refine option for Balloon tool and implement refinement process with BEN2

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 会自动

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。
 
 ## 动画
 

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 返回 |

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

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,

src/snapform/operators/balloon/__init__.py

@@ -2,4 +2,5 @@ from .operator import (
     ActivateBalloonTool,
     BalloonTool,
     ExtractBalloon,
+    RefineBalloon,
 )

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"
@@ -534,3 +637,4 @@ class BalloonTool(bpy.types.WorkSpaceTool):
         layout.prop(settings, "balloon_thickness")
         layout.prop(settings, "balloon_resolution", text="Resolution (m)")
         layout.prop(settings, "balloon_reuse_material")
+        layout.prop(settings, "balloon_ai_refine")

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(

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)

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

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
 

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