icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
马南
2026-08-20 12:17:06 +08:00
parent 4e1231790b
commit 8afcf48ddf
12 changed files with 478 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/snapform/operators/__init__.py
View File

@@ -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,

1
src/snapform/operators/balloon/__init__.py
View File

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

106
src/snapform/operators/balloon/operator.py
View File

@@ -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")
layout.prop(settings, "balloon_reuse_material")
layout.prop(settings, "balloon_ai_refine")

12
src/snapform/properties.py
View File

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

6
src/snapform/server/__init__.py
View File

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

235
src/snapform/server/jobs/balloon.py Normal file
View File

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