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

Add image extraction functionality with background removal

Browse Source 
- Implemented ExtractImageRegion operator for extracting image regions using box or lasso selection.
- Added RemoveImageRegionBackground job for background removal from selected image regions.
- Created utility functions for polygon bounds, selection canvas creation, and applying selection masks.
- Integrated image processing using PIL and numpy for efficient pixel manipulation.
- Enhanced user experience with modal drawing and status updates during extraction.
This commit is contained in:
马南
2026-08-22 15:47:34 +08:00
parent 988332548e
commit f42cba1932
27 changed files with 2513 additions and 761 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/anyimage/__init__.py
View File

@@ -6,7 +6,7 @@ from .clipboard_image import CLASSES as CLIPBOARD_CLASSES
from .runtime import runtime
from .operators import CLASSES as OPERATOR_CLASSES
from .operators.balloon import BalloonTool
from .operators.image_lasso import ImageBoxTool, ImageLassoTool
from .operators.image_extract import ImageBoxTool, ImageLassoTool, ImagePerspectiveTool
from .preferences import AnyImagePreferences
from .properties import AnyImageSettings
from .menu import AnyImageImageMenu, draw_image_context_menu
@@ -46,6 +46,11 @@ def register():
after={ImageBoxTool.bl_idname},
separator=False,
)
bpy.utils.register_tool(
ImagePerspectiveTool,
after={ImageLassoTool.bl_idname},
separator=False,
)
bpy.utils.register_tool(BalloonTool, separator=False)
bpy.types.Scene.anyimage_settings = bpy.props.PointerProperty(
type=AnyImageSettings,
@@ -60,6 +65,7 @@ def unregister():
runtime.unregister()
clipboard_image.unregister_keymaps()
bpy.utils.unregister_tool(BalloonTool)
bpy.utils.unregister_tool(ImagePerspectiveTool)
bpy.utils.unregister_tool(ImageLassoTool)
bpy.utils.unregister_tool(ImageBoxTool)
bpy.types.VIEW3D_MT_object_context_menu.remove(draw_image_context_menu)

10
src/anyimage/operators/__init__.py
View File

@@ -5,11 +5,7 @@ from .balloon import (
AnalyzeBalloon,
)
from .debug import DebugBen2, DebugDa3, DebugMoge2
from .image_lasso import (
ActivateImageLassoTool,
AnalyzeImageLasso,
ExtractImageLasso,
)
from .image_extract import CLASSES as IMAGE_EXTRACT_CLASSES
from .depth_plane import (
ConvertToDepthPlane,
ExtractOriginalFrames,
@@ -44,9 +40,7 @@ CLASSES = (
GenerateOriginalBackground,
ExtractOriginalFrames,
ClearModels,
ExtractImageLasso,
AnalyzeImageLasso,
ActivateImageLassoTool,
*IMAGE_EXTRACT_CLASSES,
ExtractBalloon,
AnalyzeBalloon,
ActivateBalloonTool,

33
src/anyimage/operators/image_extract/__init__.py Normal file
View File

@@ -0,0 +1,33 @@
from .box import ImageBoxTool
from .lasso import ImageLassoTool
from .perspective import (
ExtractImagePerspective,
ImagePerspectiveTool,
RemoveImagePerspectiveBackground,
)
from .region import (
ActivateImageExtractTool,
ExtractImageRegion,
RemoveImageRegionBackground,
)
CLASSES = (
ExtractImageRegion,
RemoveImageRegionBackground,
ActivateImageExtractTool,
ExtractImagePerspective,
RemoveImagePerspectiveBackground,
)
TOOLS = (ImageBoxTool, ImageLassoTool, ImagePerspectiveTool)
__all__ = (
"CLASSES",
"TOOLS",
"ImageBoxTool",
"ImageLassoTool",
"ImagePerspectiveTool",
)

24
src/anyimage/operators/image_extract/box.py Normal file
View File

@@ -0,0 +1,24 @@
import bpy
from .common import draw_image_extract_settings
from .region import ExtractImageRegion
class ImageBoxTool(bpy.types.WorkSpaceTool):
bl_space_type = "VIEW_3D"
bl_context_mode = "OBJECT"
bl_idname = "anyimage.image_box"
bl_label = "Image Box"
bl_description = "Extract a boxed region from an Image Empty"
bl_icon = "ops.sculpt.border_mask"
bl_operator = ExtractImageRegion.bl_idname
bl_keymap = (
(
ExtractImageRegion.bl_idname,
{"type": "LEFTMOUSE", "value": "PRESS"},
{"properties": [("selection_shape", "BOX")]},
),
)
def draw_settings(context, layout, _tool):
draw_image_extract_settings(context, layout)

349
src/anyimage/operators/image_extract/common.py Normal file
View File

@@ -0,0 +1,349 @@
import json
import math
import bpy
from ...properties import (
BEN2_MODEL_KEY,
ai_ready,
ai_setup_label,
ai_status,
image_extract_continuous_property as continuous_property,
image_extract_remove_background_property as remove_background_property,
image_extract_padding_property as padding_property,
)
from ..common import production_device, require_environment, require_model
from ..image import image_base_name, is_image_empty, plane_bounds
EXTRACT_LAYER_PROPERTY = "_anyimage_extract_layer"
EXTRACT_LAYER_STEP = 0.001
def image_extract_poll(context):
return context.area is not None and context.area.type == "VIEW_3D"
def _restore_selection(context, active_object, selected_objects):
bpy.ops.object.select_all(action="DESELECT")
for selected in selected_objects:
if bpy.data.objects.get(selected.name) is not None:
selected.select_set(True)
if active_object is not None and bpy.data.objects.get(active_object.name) is not None:
context.view_layer.objects.active = active_object
def resolve_image_extract_click(context, event, report):
"""Resolve the first click into source selection or extraction start."""
active_object = context.active_object
selected_objects = tuple(context.selected_objects)
try:
bpy.ops.view3d.select(
location=(event.mouse_region_x, event.mouse_region_y),
extend=False,
deselect_all=False,
)
except RuntimeError:
pass
clicked_object = context.active_object
if is_image_empty(clicked_object) and clicked_object != active_object:
return clicked_object, False
if is_image_empty(clicked_object):
return clicked_object, True
if is_image_empty(active_object):
_restore_selection(context, active_object, selected_objects)
return active_object, True
report({"WARNING"}, "Click an Image Empty to use an Image tool")
return None, False
def capture_view_direction(context):
from mathutils import Vector
direction = (
context.region_data.view_matrix.inverted().to_3x3()
@ Vector((0.0, 0.0, 1.0))
).normalized()
return json.dumps(tuple(direction), separators=(",", ":"))
def decode_view_direction(data):
from mathutils import Vector
direction = Vector(json.loads(data))
if direction.length <= 1e-8:
raise ValueError("The stored View direction is invalid")
return direction.normalized()
def polygon_bounds(polygon):
import numpy as np
points = np.asarray(polygon, dtype=np.float64)
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
raise ValueError("The selection does not define a valid image region")
if not np.isfinite(points).all():
raise ValueError("The selection contains invalid image coordinates")
left = int(math.floor(float(points[:, 0].min())))
top = int(math.floor(float(points[:, 1].min())))
right = int(math.ceil(float(points[:, 0].max())))
bottom = int(math.ceil(float(points[:, 1].max())))
if right <= left or bottom <= top:
raise ValueError("The selection does not define a valid image region")
return left, top, right, bottom
def crop_pixels_to_alpha(pixels, bounds, padding=0):
"""Crop Blender-order RGBA pixels to nonzero Alpha plus pixel padding."""
import numpy as np
if padding < 0:
raise ValueError("Padding must not be negative")
left, top, right, bottom = bounds
width = right - left
height = bottom - top
rgba = np.flipud(
np.asarray(pixels, dtype=np.float32).reshape((height, width, 4))
)
visible_y, visible_x = np.nonzero(rgba[:, :, 3] > 0.0)
if not len(visible_x):
raise ValueError("The selection contains no visible image pixels")
content_left = int(visible_x.min())
content_top = int(visible_y.min())
content_right = int(visible_x.max()) + 1
content_bottom = int(visible_y.max()) + 1
output_width = content_right - content_left + padding * 2
output_height = content_bottom - content_top + padding * 2
output = np.zeros((output_height, output_width, 4), dtype=np.float32)
output[
padding : padding + content_bottom - content_top,
padding : padding + content_right - content_left,
] = rgba[content_top:content_bottom, content_left:content_right]
output[output[:, :, 3] <= 0.0] = 0.0
return np.flipud(output).ravel(), (
left + content_left - padding,
top + content_top - padding,
left + content_right + padding,
top + content_bottom + padding,
)
def align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds):
from mathutils import Matrix
source_width, source_height = source_size
left, top, right, bottom = crop_bounds
x_min, x_max, y_min, y_max = plane_bounds(source_object)
local_width = x_max - x_min
local_height = y_max - y_min
crop_x_min = x_min + left * local_width / source_width
crop_x_max = x_min + right * local_width / source_width
crop_y_max = y_max - top * local_height / source_height
crop_y_min = y_max - bottom * local_height / source_height
center = (
(crop_x_min + crop_x_max) * 0.5,
(crop_y_min + crop_y_max) * 0.5,
0.0,
)
source_object.matrix_world = matrix_world @ Matrix.Translation(center)
source_object.empty_display_size = max(
crop_x_max - crop_x_min,
crop_y_max - crop_y_min,
)
source_object.empty_image_offset = (-0.5, -0.5)
def replace_empty_image_pixels(
source_object,
matrix_world,
pixels,
output_size,
placement_bounds,
suffix=".extract",
):
source_image = source_object.data
source_size = tuple(source_image.size)
output_width, output_height = (int(value) for value in output_size)
if output_width <= 0 or output_height <= 0:
raise ValueError("The extracted image dimensions must be positive")
image = bpy.data.images.new(
f"{image_base_name(source_image)}{suffix}",
width=output_width,
height=output_height,
alpha=True,
)
try:
image.colorspace_settings.name = source_image.colorspace_settings.name
image.alpha_mode = "STRAIGHT"
image.pixels.foreach_set(pixels)
image.pack()
image.update()
align_empty_to_crop(
source_object, matrix_world, source_size, placement_bounds
)
source_object.data = image
except Exception:
bpy.data.images.remove(image, do_unlink=True)
raise
return image
def replace_empty_image_from_path(source_object, matrix_world, path, crop_bounds):
source_size = tuple(source_object.data.size)
colorspace = source_object.data.colorspace_settings.name
try:
image = bpy.data.images.load(str(path), check_existing=False)
except (OSError, RuntimeError) as error:
raise RuntimeError("Unable to load the extracted image") from error
try:
image.name = f"{image_base_name(source_object.data)}.extract"
image.colorspace_settings.name = colorspace
image.alpha_mode = "STRAIGHT"
image.pack()
align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds)
source_object.data = image
except Exception:
bpy.data.images.remove(image, do_unlink=True)
raise
return image
def image_extract_target(context, source_object, continuous):
if not continuous:
return source_object, 0
layer = int(source_object.get(EXTRACT_LAYER_PROPERTY, 0)) + 1
source_object[EXTRACT_LAYER_PROPERTY] = layer
target = source_object.copy()
target[EXTRACT_LAYER_PROPERTY] = layer
collection = (
source_object.users_collection[0]
if source_object.users_collection
else context.collection
)
collection.objects.link(target)
return target, layer
def offset_continuous_result(target, view_direction_data, layer):
if layer <= 0:
return
direction = decode_view_direction(view_direction_data)
world_scale = max(abs(value) for value in target.matrix_world.to_scale())
distance = target.empty_display_size * world_scale * EXTRACT_LAYER_STEP * layer
target.matrix_world.translation += direction * distance
def select_image_extract_result(context, source_object, target, continuous):
bpy.ops.object.select_all(action="DESELECT")
selected = source_object if continuous else target
selected.select_set(True)
context.view_layer.objects.active = selected
def image_extract_background_request(input_path, selection_data, padding):
require_environment()
require_model(BEN2_MODEL_KEY)
polygon = json.loads(selection_data)
polygon_bounds(polygon)
return {
"input": str(input_path),
"polygon": polygon,
"padding": padding,
"model": BEN2_MODEL_KEY,
"device": production_device().lower(),
}
def apply_image_extract_settings(operator, context):
settings = context.scene.anyimage_settings
operator.remove_background = (
ai_ready() and settings.image_extract_remove_background
)
operator.continuous = settings.image_extract_continuous
operator.padding = settings.image_extract_padding
def draw_image_extract_operator_settings(layout, operator):
status = ai_status()
row = layout.row()
row.enabled = status["ready"]
row.prop(operator, "remove_background")
layout.prop(operator, "continuous")
layout.prop(operator, "padding")
if not status["ready"]:
layout.operator(
"anyimage.setup_ai_environment",
text=ai_setup_label(status),
icon="IMPORT",
)
def draw_image_extract_settings(context, layout):
settings = context.scene.anyimage_settings
status = ai_status()
row = layout.row()
row.enabled = status["ready"]
row.prop(settings, "image_extract_remove_background")
layout.prop(settings, "image_extract_continuous")
layout.prop(settings, "image_extract_padding")
if not status["ready"]:
layout.operator(
"anyimage.setup_ai_environment",
text=ai_setup_label(status),
icon="IMPORT",
)
def dashed_line_segments(vertices, dash_length=4.0, gap_length=4.0):
if dash_length <= 0.0 or gap_length < 0.0:
raise ValueError("Dash length must be positive and gap length cannot be negative")
pattern_length = dash_length + gap_length
segments = []
distance_offset = 0.0
for start, end in zip(vertices, vertices[1:]):
start_x, start_y = (float(value) for value in start)
end_x, end_y = (float(value) for value in end)
segment_length = math.hypot(end_x - start_x, end_y - start_y)
if segment_length <= 0.0:
continue
direction_x = (end_x - start_x) / segment_length
direction_y = (end_y - start_y) / segment_length
position = 0.0
while position < segment_length:
phase = (distance_offset + position) % pattern_length
if phase < dash_length:
run_length = min(dash_length - phase, segment_length - position)
segments.extend(
(
(
start_x + direction_x * position,
start_y + direction_y * position,
),
(
start_x + direction_x * (position + run_length),
start_y + direction_y * (position + run_length),
),
)
)
else:
run_length = min(pattern_length - phase, segment_length - position)
position += run_length
distance_offset += segment_length
return tuple(segments)
def draw_dashed_line(vertices, color=(1.0, 1.0, 1.0, 0.9)):
import gpu
from gpu_extras.batch import batch_for_shader
segments = dashed_line_segments(vertices)
if not segments:
return
shader = gpu.shader.from_builtin("UNIFORM_COLOR")
batch = batch_for_shader(shader, "LINES", {"pos": segments})
gpu.state.blend_set("ALPHA")
gpu.state.line_width_set(1.0)
shader.bind()
shader.uniform_float("color", color)
batch.draw(shader)
gpu.state.blend_set("NONE")

24
src/anyimage/operators/image_extract/lasso.py Normal file
View File

@@ -0,0 +1,24 @@
import bpy
from .common import draw_image_extract_settings
from .region import ExtractImageRegion
class ImageLassoTool(bpy.types.WorkSpaceTool):
bl_space_type = "VIEW_3D"
bl_context_mode = "OBJECT"
bl_idname = "anyimage.image_lasso"
bl_label = "Image Lasso"
bl_description = "Extract a lassoed region from an Image Empty"
bl_icon = "ops.sculpt.lasso_mask"
bl_operator = ExtractImageRegion.bl_idname
bl_keymap = (
(
ExtractImageRegion.bl_idname,
{"type": "LEFTMOUSE", "value": "PRESS"},
{"properties": [("selection_shape", "LASSO")]},
),
)
def draw_settings(context, layout, _tool):
draw_image_extract_settings(context, layout)

1020
src/anyimage/operators/image_extract/perspective.py Normal file
View File

File diff suppressed because it is too large Load Diff

378
src/anyimage/operators/image_extract/region.py Normal file
View File

@@ -0,0 +1,378 @@
import json
import bpy
from ...runtime import JobOperatorBase
from ..balloon.operator import (
deserialize_matrix,
drawing_in_region,
hidden_string_property,
screen_path_to_image_pixels,
serialize_matrix,
simplify_screen_path,
)
from ..common import HIDDEN, cleanup_temporary_input, push_undo
from ..image import (
is_image_empty,
plane_bounds,
prepare_image_input,
remove_temporary_input,
video_total_frames,
)
from .common import (
apply_image_extract_settings,
capture_view_direction,
continuous_property,
image_extract_background_request,
crop_pixels_to_alpha,
draw_dashed_line,
draw_image_extract_operator_settings,
remove_background_property,
resolve_image_extract_click,
image_extract_poll,
image_extract_target,
offset_continuous_result,
padding_property,
polygon_bounds,
replace_empty_image_from_path,
replace_empty_image_pixels,
select_image_extract_result,
)
def points_inside_polygon(points, polygon):
import numpy as np
points = np.asarray(points, dtype=np.float64)
polygon = np.asarray(polygon, dtype=np.float64)
inside = np.zeros(len(points), dtype=bool)
previous = polygon[-1]
for current in polygon:
crosses = (current[1] > points[:, 1]) != (previous[1] > points[:, 1])
edge_x = (
(previous[0] - current[0])
* (points[:, 1] - current[1])
/ (previous[1] - current[1] + 1e-300)
+ current[0]
)
inside ^= crosses & (points[:, 0] < edge_x)
previous = current
return inside
def extract_region_pixels(source_pixels, source_size, polygon):
"""Return Blender-order RGBA pixels and the unbounded selection rectangle."""
import numpy as np
source_width, source_height = (int(value) for value in source_size)
if source_width <= 0 or source_height <= 0:
raise ValueError("The source image has no readable dimensions")
source = np.asarray(source_pixels, dtype=np.float32)
if source.size != source_width * source_height * 4:
raise ValueError("The source image pixels do not match its dimensions")
source = np.flipud(source.reshape((source_height, source_width, 4)))
left, top, right, bottom = polygon_bounds(polygon)
width = right - left
height = bottom - top
result = np.zeros((height, width, 4), dtype=np.float32)
source_left = max(left, 0)
source_top = max(top, 0)
source_right = min(right, source_width)
source_bottom = min(bottom, source_height)
if source_right <= source_left or source_bottom <= source_top:
raise ValueError("The selection does not overlap the source image")
result[
source_top - top : source_bottom - top,
source_left - left : source_right - left,
] = source[source_top:source_bottom, source_left:source_right]
y, x = np.indices((height, width), dtype=np.float64)
centers = np.column_stack(
((x.ravel() + left + 0.5), (y.ravel() + top + 0.5))
)
selected = points_inside_polygon(centers, polygon).reshape((height, width))
result[~selected] = 0.0
result[result[:, :, 3] <= 0.0] = 0.0
return np.flipud(result).ravel(), (left, top, right, bottom)
def selection_shape_property():
return bpy.props.EnumProperty(
items=(
("LASSO", "Lasso", "Draw a freehand image extraction region"),
("BOX", "Box", "Draw a rectangular image extraction region"),
),
default="BOX",
options={"HIDDEN", "SKIP_SAVE"},
)
def box_screen_path(path):
if len(path) < 2:
return ()
x0, y0 = path[0]
x1, y1 = path[-1]
if abs(x1 - x0) < 2.0 or abs(y1 - y0) < 2.0:
return ()
return ((x0, y0), (x1, y0), (x1, y1), (x0, y1))
class ExtractImageRegion(bpy.types.Operator):
bl_idname = "anyimage.extract_image_region"
bl_label = "Extract Image Region"
bl_description = "Extract a boxed or lassoed region from an Image Empty"
bl_options = {"REGISTER", "UNDO"}
source_object_name: hidden_string_property()
source_matrix_data: hidden_string_property()
view_direction_data: hidden_string_property()
selection_data: hidden_string_property()
remove_background: remove_background_property()
continuous: continuous_property()
padding: padding_property()
selection_shape: selection_shape_property()
_path = None
_handle = None
_area_pointer = 0
_region_pointer = 0
def draw(self, _context):
draw_image_extract_operator_settings(self.layout, self)
@classmethod
def poll(cls, context):
return image_extract_poll(context)
def _draw_selection(self):
if not self._path or not drawing_in_region(
self._area_pointer, self._region_pointer
):
return
vertices = (
list(box_screen_path(self._path))
if self.selection_shape == "BOX"
else list(self._path)
)
if len(vertices) > 2:
vertices.append(vertices[0])
draw_dashed_line(vertices)
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, should_extract = resolve_image_extract_click(
context, event, self.report
)
if not should_extract:
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_image_extract_settings(self, context)
self.source_object_name = source_object.name
self.source_matrix_data = serialize_matrix(source_object.matrix_world)
self.view_direction_data = capture_view_direction(context)
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_selection, (), "WINDOW", "POST_PIXEL"
)
context.workspace.status_text_set(
f"Image {self.selection_shape.title()}: 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":
self._path.append((event.mouse_region_x, event.mouse_region_y))
path = (
list(box_screen_path(self._path))
if self.selection_shape == "BOX"
else simplify_screen_path(self._path)
)
if len(path) < 3:
self.report({"WARNING"}, "Draw a larger image region")
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:
polygon = screen_path_to_image_pixels(
path,
context.region,
context.region_data,
deserialize_matrix(self.source_matrix_data),
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)
polygon_bounds(polygon)
if video_total_frames(source_object.data) > 1:
raise RuntimeError("Image tools only support a still image")
if self.remove_background:
source_path, delete_input = prepare_image_input(
source_object.data, context.scene
)
result = bpy.ops.anyimage.remove_image_region_background(
input_path=str(source_path),
delete_input=delete_input,
source_object_name=self.source_object_name,
source_matrix_data=self.source_matrix_data,
view_direction_data=self.view_direction_data,
selection_data=self.selection_data,
continuous=self.continuous,
padding=self.padding,
)
if "CANCELLED" in result:
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
return {"FINISHED"}
import numpy as np
image = source_object.data
source_pixels = np.empty(len(image.pixels), dtype=np.float32)
image.pixels.foreach_get(source_pixels)
pixels, crop_bounds = extract_region_pixels(
source_pixels, tuple(image.size), polygon
)
pixels, crop_bounds = crop_pixels_to_alpha(
pixels, crop_bounds, self.padding
)
target, layer = image_extract_target(
context, source_object, self.continuous
)
try:
left, top, right, bottom = crop_bounds
replace_empty_image_pixels(
target,
deserialize_matrix(self.source_matrix_data),
pixels,
(right - left, bottom - top),
crop_bounds,
)
offset_continuous_result(
target, self.view_direction_data, layer
)
except Exception:
if target is not source_object:
bpy.data.objects.remove(target, do_unlink=True)
raise
select_image_extract_result(
context, source_object, target, self.continuous
)
except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class RemoveImageRegionBackground(JobOperatorBase, bpy.types.Operator):
bl_idname = "anyimage.remove_image_region_background"
bl_label = "Remove Image Region Background"
job_type = "extract-image-background"
starting_message = "Removing image region background"
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()
view_direction_data: hidden_string_property()
selection_data: hidden_string_property()
continuous: continuous_property()
padding: padding_property()
def request(self, _context):
return image_extract_background_request(
self.input_path, self.selection_data, self.padding
)
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")
crop_bounds = tuple(
json.loads(result.file("bounds").read_text(encoding="utf-8"))
)
target, layer = image_extract_target(
context, source_object, self.continuous
)
try:
replace_empty_image_from_path(
target,
deserialize_matrix(self.source_matrix_data),
result.file("image"),
crop_bounds,
)
offset_continuous_result(target, self.view_direction_data, layer)
except Exception:
if target is not source_object:
bpy.data.objects.remove(target, do_unlink=True)
raise
select_image_extract_result(
context, source_object, target, self.continuous
)
push_undo("Extract Image Region")
return "Image extraction with background removal is ready"
def cleanup(self):
cleanup_temporary_input(self.input_path, self.delete_input)
class ActivateImageExtractTool(bpy.types.Operator):
bl_idname = "anyimage.activate_image_extract_tool"
bl_label = "Image Extract Tool"
bl_description = "Activate the Image Box tool in the 3D Viewport"
bl_options = {"INTERNAL"}
@classmethod
def poll(cls, context):
return image_extract_poll(context)
def execute(self, _context):
return bpy.ops.wm.tool_set_by_id(
name="anyimage.image_box", space_type="VIEW_3D"
)

635
src/anyimage/operators/image_lasso.py
View File

@@ -1,635 +0,0 @@
import json
import math
import bpy
from ..properties import (
BEN2_MODEL_KEY,
ai_ready,
ai_setup_label,
ai_status,
image_lasso_contour_analysis_property as contour_analysis_property,
image_lasso_keep_original_property as keep_original_property,
image_lasso_padding_property as padding_property,
)
from ..runtime import JobOperatorBase
from .balloon.operator import (
deserialize_matrix,
drawing_in_region,
hidden_string_property,
screen_path_to_image_pixels,
serialize_matrix,
simplify_screen_path,
)
from .common import (
HIDDEN,
cleanup_temporary_input,
production_device,
push_undo,
require_environment,
require_model,
)
from .image import (
image_base_name,
is_image_empty,
plane_bounds,
prepare_image_input,
remove_temporary_input,
video_total_frames,
)
def lasso_bounds(polygon):
import numpy as np
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")
left = int(math.floor(float(points[:, 0].min())))
top = int(math.floor(float(points[:, 1].min())))
right = int(math.ceil(float(points[:, 0].max())))
bottom = int(math.ceil(float(points[:, 1].max())))
if right <= left or bottom <= top:
raise ValueError("The lasso does not define a valid image region")
return left, top, right, bottom
def points_inside_polygon(points, polygon):
import numpy as np
points = np.asarray(points, dtype=np.float64)
polygon = np.asarray(polygon, dtype=np.float64)
inside = np.zeros(len(points), dtype=bool)
previous = polygon[-1]
for current in polygon:
crosses = (current[1] > points[:, 1]) != (previous[1] > points[:, 1])
edge_x = (
(previous[0] - current[0])
* (points[:, 1] - current[1])
/ (previous[1] - current[1] + 1e-300)
+ current[0]
)
inside ^= crosses & (points[:, 0] < edge_x)
previous = current
return inside
def extract_lasso_pixels(source_pixels, source_size, polygon):
"""Return Blender-order RGBA pixels and the unbounded lasso rectangle."""
import numpy as np
source_width, source_height = (int(value) for value in source_size)
if source_width <= 0 or source_height <= 0:
raise ValueError("The source image has no readable dimensions")
source = np.asarray(source_pixels, dtype=np.float32)
if source.size != source_width * source_height * 4:
raise ValueError("The source image pixels do not match its dimensions")
source = np.flipud(source.reshape((source_height, source_width, 4)))
left, top, right, bottom = lasso_bounds(polygon)
width = right - left
height = bottom - top
result = np.zeros((height, width, 4), dtype=np.float32)
source_left = max(left, 0)
source_top = max(top, 0)
source_right = min(right, source_width)
source_bottom = min(bottom, source_height)
if source_right <= source_left or source_bottom <= source_top:
raise ValueError("The lasso does not overlap the source image")
result[
source_top - top : source_bottom - top,
source_left - left : source_right - left,
] = source[source_top:source_bottom, source_left:source_right]
y, x = np.indices((height, width), dtype=np.float64)
centers = np.column_stack(
((x.ravel() + left + 0.5), (y.ravel() + top + 0.5))
)
selected = points_inside_polygon(centers, polygon).reshape((height, width))
result[~selected] = 0.0
result[result[:, :, 3] <= 0.0] = 0.0
return np.flipud(result).ravel(), (left, top, right, bottom)
def crop_pixels_to_alpha(pixels, bounds, padding=0):
"""Crop Blender-order RGBA pixels to nonzero Alpha plus pixel padding."""
import numpy as np
if padding < 0:
raise ValueError("Padding must not be negative")
left, top, right, bottom = bounds
width = right - left
height = bottom - top
rgba = np.flipud(
np.asarray(pixels, dtype=np.float32).reshape((height, width, 4))
)
visible_y, visible_x = np.nonzero(rgba[:, :, 3] > 0.0)
if not len(visible_x):
raise ValueError("The lasso contains no visible image pixels")
content_left = int(visible_x.min())
content_top = int(visible_y.min())
content_right = int(visible_x.max()) + 1
content_bottom = int(visible_y.max()) + 1
output_width = content_right - content_left + padding * 2
output_height = content_bottom - content_top + padding * 2
output = np.zeros((output_height, output_width, 4), dtype=np.float32)
output[
padding : padding + content_bottom - content_top,
padding : padding + content_right - content_left,
] = rgba[content_top:content_bottom, content_left:content_right]
output[output[:, :, 3] <= 0.0] = 0.0
return np.flipud(output).ravel(), (
left + content_left - padding,
top + content_top - padding,
left + content_right + padding,
top + content_bottom + padding,
)
def align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds):
from mathutils import Matrix
source_width, source_height = source_size
left, top, right, bottom = crop_bounds
x_min, x_max, y_min, y_max = plane_bounds(source_object)
local_width = x_max - x_min
local_height = y_max - y_min
crop_x_min = x_min + left * local_width / source_width
crop_x_max = x_min + right * local_width / source_width
crop_y_max = y_max - top * local_height / source_height
crop_y_min = y_max - bottom * local_height / source_height
center = (
(crop_x_min + crop_x_max) * 0.5,
(crop_y_min + crop_y_max) * 0.5,
0.0,
)
source_object.matrix_world = matrix_world @ Matrix.Translation(center)
source_object.empty_display_size = max(
crop_x_max - crop_x_min,
crop_y_max - crop_y_min,
)
source_object.empty_image_offset = (-0.5, -0.5)
def replace_empty_image(source_object, matrix_world, pixels, crop_bounds):
source_image = source_object.data
source_size = tuple(source_image.size)
left, top, right, bottom = crop_bounds
image = bpy.data.images.new(
f"{image_base_name(source_image)}.extract",
width=right - left,
height=bottom - top,
alpha=True,
)
try:
image.colorspace_settings.name = source_image.colorspace_settings.name
image.alpha_mode = "STRAIGHT"
image.pixels.foreach_set(pixels)
image.pack()
image.update()
align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds)
source_object.data = image
except Exception:
bpy.data.images.remove(image, do_unlink=True)
raise
return image
def replace_empty_image_from_path(source_object, matrix_world, path, crop_bounds):
source_size = tuple(source_object.data.size)
colorspace = source_object.data.colorspace_settings.name
try:
image = bpy.data.images.load(str(path), check_existing=False)
except (OSError, RuntimeError) as error:
raise RuntimeError("Unable to load the extracted image") from error
try:
image.name = f"{image_base_name(source_object.data)}.extract"
image.colorspace_settings.name = colorspace
image.alpha_mode = "STRAIGHT"
image.pack()
align_empty_to_crop(source_object, matrix_world, source_size, crop_bounds)
source_object.data = image
except Exception:
bpy.data.images.remove(image, do_unlink=True)
raise
return image
def image_lasso_target(context, source_object, keep_original):
if not keep_original:
return source_object
target = source_object.copy()
collection = (
source_object.users_collection[0]
if source_object.users_collection
else context.collection
)
collection.objects.link(target)
return target
def select_image_lasso_result(context, target):
bpy.ops.object.select_all(action="DESELECT")
target.select_set(True)
context.view_layer.objects.active = target
def selection_shape_property():
return bpy.props.EnumProperty(
items=(
("LASSO", "Lasso", "Draw a freehand image extraction region"),
("BOX", "Box", "Draw a rectangular image extraction region"),
),
default="BOX",
options={"HIDDEN", "SKIP_SAVE"},
)
def box_screen_path(path):
if len(path) < 2:
return ()
x0, y0 = path[0]
x1, y1 = path[-1]
if abs(x1 - x0) < 2.0 or abs(y1 - y0) < 2.0:
return ()
return ((x0, y0), (x1, y0), (x1, y1), (x0, y1))
def dashed_line_segments(vertices, dash_length=4.0, gap_length=4.0):
if dash_length <= 0.0 or gap_length < 0.0:
raise ValueError("Dash length must be positive and gap length cannot be negative")
pattern_length = dash_length + gap_length
segments = []
distance_offset = 0.0
for start, end in zip(vertices, vertices[1:]):
start_x, start_y = (float(value) for value in start)
end_x, end_y = (float(value) for value in end)
segment_length = math.hypot(end_x - start_x, end_y - start_y)
if segment_length <= 0.0:
continue
direction_x = (end_x - start_x) / segment_length
direction_y = (end_y - start_y) / segment_length
position = 0.0
while position < segment_length:
phase = (distance_offset + position) % pattern_length
if phase < dash_length:
run_length = min(dash_length - phase, segment_length - position)
segments.extend(
(
(
start_x + direction_x * position,
start_y + direction_y * position,
),
(
start_x + direction_x * (position + run_length),
start_y + direction_y * (position + run_length),
),
)
)
else:
run_length = min(pattern_length - phase, segment_length - position)
position += run_length
distance_offset += segment_length
return tuple(segments)
def draw_dashed_line(vertices):
import gpu
from gpu_extras.batch import batch_for_shader
segments = dashed_line_segments(vertices)
if not segments:
return
shader = gpu.shader.from_builtin("UNIFORM_COLOR")
batch = batch_for_shader(shader, "LINES", {"pos": segments})
gpu.state.blend_set("ALPHA")
gpu.state.line_width_set(1.0)
shader.bind()
shader.uniform_float("color", (1.0, 1.0, 1.0, 0.9))
batch.draw(shader)
gpu.state.blend_set("NONE")
class ExtractImageLasso(bpy.types.Operator):
bl_idname = "anyimage.extract_image_lasso"
bl_label = "Extract Image"
bl_description = "Replace an Image Empty with its lassoed image region"
bl_options = {"REGISTER", "UNDO"}
source_object_name: hidden_string_property()
source_matrix_data: hidden_string_property()
selection_data: hidden_string_property()
contour_analysis: contour_analysis_property()
keep_original: keep_original_property()
padding: padding_property()
selection_shape: selection_shape_property()
_path = None
_handle = None
_area_pointer = 0
_region_pointer = 0
def draw(self, _context):
status = ai_status()
row = self.layout.row()
row.enabled = status["ready"]
row.prop(self, "contour_analysis")
self.layout.prop(self, "keep_original")
self.layout.prop(self, "padding")
if not status["ready"]:
self.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_selection(self):
if not self._path or not drawing_in_region(
self._area_pointer, self._region_pointer
):
return
vertices = (
list(box_screen_path(self._path))
if self.selection_shape == "BOX"
else list(self._path)
)
if len(vertices) > 2:
vertices.append(vertices[0])
draw_dashed_line(vertices)
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
self.contour_analysis = (
ai_ready() and context.scene.anyimage_settings.image_lasso_contour_analysis
)
self.keep_original = context.scene.anyimage_settings.image_lasso_keep_original
self.padding = context.scene.anyimage_settings.image_lasso_padding
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_selection, (), "WINDOW", "POST_PIXEL"
)
context.workspace.status_text_set(
f"Image {self.selection_shape.title()}: 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":
self._path.append((event.mouse_region_x, event.mouse_region_y))
path = (
list(box_screen_path(self._path))
if self.selection_shape == "BOX"
else simplify_screen_path(self._path)
)
if len(path) < 3:
self.report({"WARNING"}, "Draw a larger image region")
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)
crop_bounds = lasso_bounds(polygon)
if video_total_frames(source_object.data) > 1:
raise RuntimeError("Image Lasso only supports a still image")
if self.contour_analysis:
source_path, delete_input = prepare_image_input(
source_object.data, context.scene
)
result = bpy.ops.anyimage.analyze_image_lasso(
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,
keep_original=self.keep_original,
padding=self.padding,
)
if "CANCELLED" in result:
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
return {"FINISHED"}
import numpy as np
image = source_object.data
source_pixels = np.empty(len(image.pixels), dtype=np.float32)
image.pixels.foreach_get(source_pixels)
pixels, crop_bounds = extract_lasso_pixels(
source_pixels, tuple(image.size), polygon
)
pixels, crop_bounds = crop_pixels_to_alpha(
pixels, crop_bounds, self.padding
)
target = image_lasso_target(context, source_object, self.keep_original)
try:
replace_empty_image(
target,
deserialize_matrix(self.source_matrix_data),
pixels,
crop_bounds,
)
except Exception:
if target is not source_object:
bpy.data.objects.remove(target, do_unlink=True)
raise
select_image_lasso_result(
context,
source_object if self.keep_original else target,
)
except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class AnalyzeImageLasso(JobOperatorBase, bpy.types.Operator):
bl_idname = "anyimage.analyze_image_lasso"
bl_label = "Analyze Image Lasso"
job_type = "analyze-image-lasso"
starting_message = "Analyzing image contour"
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()
keep_original: keep_original_property()
padding: padding_property()
def request(self, _context):
require_environment()
require_model(BEN2_MODEL_KEY)
polygon = json.loads(self.selection_data)
lasso_bounds(polygon)
return {
"input": str(self.input_path),
"polygon": polygon,
"padding": self.padding,
"model": BEN2_MODEL_KEY,
"device": production_device().lower(),
}
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")
crop_bounds = tuple(
json.loads(result.file("bounds").read_text(encoding="utf-8"))
)
target = image_lasso_target(context, source_object, self.keep_original)
try:
replace_empty_image_from_path(
target,
deserialize_matrix(self.source_matrix_data),
result.file("image"),
crop_bounds,
)
except Exception:
if target is not source_object:
bpy.data.objects.remove(target, do_unlink=True)
raise
select_image_lasso_result(
context,
source_object if self.keep_original else target,
)
push_undo("Extract Image")
return "Contoured image extraction is ready"
def cleanup(self):
cleanup_temporary_input(self.input_path, self.delete_input)
class ActivateImageLassoTool(bpy.types.Operator):
bl_idname = "anyimage.activate_image_lasso_tool"
bl_label = "Image Extract Tool"
bl_description = "Activate the Image Box tool in the 3D Viewport"
bl_options = {"INTERNAL"}
@classmethod
def poll(cls, context):
return ExtractImageLasso.poll(context)
def execute(self, _context):
return bpy.ops.wm.tool_set_by_id(
name=ImageBoxTool.bl_idname, space_type="VIEW_3D"
)
class ImageBoxTool(bpy.types.WorkSpaceTool):
bl_space_type = "VIEW_3D"
bl_context_mode = "OBJECT"
bl_idname = "anyimage.image_box"
bl_label = "Image Box"
bl_description = "Replace an Image Empty with one boxed image region"
bl_icon = "ops.sculpt.border_mask"
bl_operator = ExtractImageLasso.bl_idname
bl_keymap = (
(
ExtractImageLasso.bl_idname,
{"type": "LEFTMOUSE", "value": "PRESS"},
{"properties": [("selection_shape", "BOX")]},
),
)
def draw_settings(context, layout, _tool):
draw_image_extract_settings(context, layout)
class ImageLassoTool(bpy.types.WorkSpaceTool):
bl_space_type = "VIEW_3D"
bl_context_mode = "OBJECT"
bl_idname = "anyimage.image_lasso"
bl_label = "Image Lasso"
bl_description = "Replace an Image Empty with one lassoed image region"
bl_icon = "ops.sculpt.lasso_mask"
bl_operator = ExtractImageLasso.bl_idname
bl_keymap = (
(
ExtractImageLasso.bl_idname,
{"type": "LEFTMOUSE", "value": "PRESS"},
{"properties": [("selection_shape", "LASSO")]},
),
)
def draw_settings(context, layout, _tool):
draw_image_extract_settings(context, layout)
def draw_image_extract_settings(context, layout):
settings = context.scene.anyimage_settings
status = ai_status()
row = layout.row()
row.enabled = status["ready"]
row.prop(settings, "image_lasso_contour_analysis")
layout.prop(settings, "image_lasso_keep_original")
layout.prop(settings, "image_lasso_padding")
if not status["ready"]:
layout.operator(
"anyimage.setup_ai_environment",
text=ai_setup_label(status),
icon="IMPORT",
)

22
src/anyimage/properties.py
View File

@@ -280,25 +280,27 @@ def balloon_contour_analysis_property():
)
def image_lasso_contour_analysis_property():
def image_extract_remove_background_property():
return bpy.props.BoolProperty(
name="Contour",
name="Remove Background",
description=(
"Use BEN2 to refine the lasso Alpha while preserving the source colors"
"Use BEN2 to remove the background inside the selected image region"
),
default=False,
)
def image_lasso_keep_original_property():
def image_extract_continuous_property():
return bpy.props.BoolProperty(
name="Keep Original",
description="Keep the source Image Empty and create a cropped copy",
name="Continuous",
description=(
"Keep the source Image Empty active and continuously create extracted copies"
),
default=False,
)
def image_lasso_padding_property():
def image_extract_padding_property():
return bpy.props.IntProperty(
name="Padding",
description="Transparent padding in pixels around the visible Alpha bounds",
@@ -335,9 +337,9 @@ def balloon_radial_zoom_property(scale):
class AnyImageSettings(bpy.types.PropertyGroup):
image_lasso_contour_analysis: image_lasso_contour_analysis_property()
image_lasso_keep_original: image_lasso_keep_original_property()
image_lasso_padding: image_lasso_padding_property()
image_extract_remove_background: image_extract_remove_background_property()
image_extract_continuous: image_extract_continuous_property()
image_extract_padding: image_extract_padding_property()
balloon_mode: balloon_mode_property()
balloon_thickness: balloon_thickness_property()
balloon_resolution: balloon_resolution_property()

8
src/anyimage/server/__init__.py
View File

@@ -10,7 +10,7 @@ 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
from .jobs.frames import run as _extract_frames
from .jobs.image_lasso import run as _analyze_image_lasso
from .jobs.image_extract import run as _extract_image_background
from .model_manager import ModelManager, ModelSessionCache
@@ -39,9 +39,9 @@ def analyze_balloon(context, parameters):
return _analyze_balloon(context, parameters)
@server.job("analyze-image-lasso")
def analyze_image_lasso(context, parameters):
return _analyze_image_lasso(context, parameters)
@server.job("extract-image-background")
def extract_image_background(context, parameters):
return _extract_image_background(context, parameters)
@server.job("remove-background")

21
src/anyimage/server/jobs/background.py
View File

@@ -5,6 +5,20 @@ from ..media import prepare_background_input
from ..models.ben2 import infer_foreground, infer_foreground_frames
def remove_image_background(
input_path, model_directory, device, model_manager, cancel_check
):
"""Remove one image background with the shared BEN2 configuration."""
return infer_foreground(
input_path,
model_directory,
device,
model_cache=model_manager,
cancel_check=cancel_check,
preserve_input_alpha=True,
)
def run(context, parameters):
"""Run BEN2 and write transparent foreground images into the Job directory."""
total_started = time.perf_counter()
@@ -30,13 +44,12 @@ def run(context, parameters):
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
context.progress(0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
foreground, load_ms, inference_ms = remove_image_background(
images[0],
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
preserve_input_alpha=True,
model_manager,
context.check_cancelled,
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")

81
src/anyimage/server/jobs/image_lasso.py → src/anyimage/server/jobs/image_extract.py
View File

@@ -2,36 +2,36 @@ import json
import math
import time
from ..models.ben2 import infer_foreground
from .background import remove_image_background
def lasso_bounds(polygon):
def polygon_bounds(polygon):
import numpy as np
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")
raise ValueError("The selection does not define a valid image region")
if not np.isfinite(points).all():
raise ValueError("The lasso contains invalid image coordinates")
raise ValueError("The selection contains invalid image coordinates")
left = int(math.floor(float(points[:, 0].min())))
top = int(math.floor(float(points[:, 1].min())))
right = int(math.ceil(float(points[:, 0].max())))
bottom = int(math.ceil(float(points[:, 1].max())))
if right <= left or bottom <= top:
raise ValueError("The lasso does not define a valid image region")
raise ValueError("The selection does not define a valid image region")
return left, top, right, bottom
def create_lasso_canvas(input_path, polygon):
"""Create an unbounded lasso canvas with transparent black outside."""
def create_selection_canvas(input_path, polygon):
"""Create an unbounded selection canvas with transparent black outside."""
import numpy as np
from PIL import Image, ImageDraw
with Image.open(input_path) as opened_image:
source = opened_image.convert("RGBA")
left, top, right, bottom = lasso_bounds(polygon)
left, top, right, bottom = polygon_bounds(polygon)
if right <= 0 or bottom <= 0 or left >= source.width or top >= source.height:
raise ValueError("The lasso does not overlap the source image")
raise ValueError("The selection does not overlap the source image")
canvas = Image.new("RGBA", (right - left, bottom - top), (0, 0, 0, 0))
canvas.paste(source, (-left, -top))
relative = [(float(x) - left, float(y) - top) for x, y in polygon]
@@ -48,27 +48,23 @@ def create_lasso_canvas(input_path, polygon):
)
def apply_contour_alpha(canvas, selection_mask, foreground):
"""Keep source RGB and combine source, lasso, and BEN2 alpha."""
def apply_selection_mask(foreground, selection_mask):
"""Clear pixels outside the source selection after background removal."""
import numpy as np
from PIL import Image
source = np.asarray(canvas.convert("RGBA"), dtype=np.uint8).copy()
alpha = np.asarray(foreground.convert("RGBA").getchannel("A"), dtype=np.uint8)
if alpha.shape != selection_mask.shape:
alpha = np.asarray(
Image.fromarray(alpha, mode="L").resize(
canvas.size, Image.Resampling.BILINEAR
pixels = np.asarray(foreground.convert("RGBA"), dtype=np.uint8).copy()
if pixels.shape[:2] != selection_mask.shape:
pixels = np.asarray(
foreground.convert("RGBA").resize(
(selection_mask.shape[1], selection_mask.shape[0]),
Image.Resampling.BILINEAR,
),
dtype=np.uint8,
)
source[:, :, 3] = np.where(
selection_mask,
np.minimum(source[:, :, 3], alpha),
0,
).astype(np.uint8)
source[source[:, :, 3] == 0, :3] = 0
return Image.fromarray(source, mode="RGBA")
pixels[~selection_mask] = 0
pixels[pixels[:, :, 3] == 0, :3] = 0
return Image.fromarray(pixels, mode="RGBA")
def crop_to_alpha(image, bounds, padding=0):
@@ -81,7 +77,7 @@ def crop_to_alpha(image, bounds, padding=0):
pixels = np.asarray(image.convert("RGBA"), dtype=np.uint8)
visible_y, visible_x = np.nonzero(pixels[:, :, 3] > 0)
if not len(visible_x):
raise ValueError("The lasso contains no visible image pixels")
raise ValueError("The selection contains no visible image pixels")
content_left = int(visible_x.min())
content_top = int(visible_y.min())
content_right = int(visible_x.max()) + 1
@@ -105,47 +101,38 @@ def crop_to_alpha(image, bounds, padding=0):
def run(context, parameters):
"""Use BEN2 to refine an Image Lasso without changing its canvas bounds."""
"""Remove the background from an image selection and return source bounds."""
started = time.perf_counter()
context.progress(0.05, "Preparing image lasso")
context.progress(0.05, "Preparing image selection")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
canvas, selection_mask, bounds = create_lasso_canvas(
canvas, selection_mask, bounds = create_selection_canvas(
parameters["input"], parameters["polygon"]
)
lasso_path = output_directory / "lasso-input.png"
import numpy as np
from PIL import Image
analysis_pixels = np.asarray(canvas, dtype=np.uint8).copy()
analysis_pixels[~selection_mask] = 0
Image.fromarray(analysis_pixels, mode="RGBA").save(
lasso_path, format="PNG", compress_level=1
)
selection_path = output_directory / "selection-input.png"
canvas.save(selection_path, format="PNG", compress_level=1)
model_manager = context.resource("model_manager")
context.progress(0.2, f"Loading BEN2 for {parameters['device'].upper()}")
foreground, _load_ms, _inference_ms = infer_foreground(
lasso_path,
foreground, _load_ms, _inference_ms = remove_image_background(
selection_path,
model_manager.directory(parameters["model"]),
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
model_manager,
context.check_cancelled,
)
context.check_cancelled()
context.progress(0.9, "Applying contour alpha")
context.progress(0.9, "Applying selection mask")
output, bounds = crop_to_alpha(
apply_contour_alpha(canvas, selection_mask, foreground),
apply_selection_mask(foreground, selection_mask),
bounds,
int(parameters.get("padding", 0)),
)
output_path = output_directory / "image.png"
output.save(
output_path, format="PNG", compress_level=1
)
output.save(output_path, format="PNG", compress_level=1)
bounds_path = output_directory / "bounds.json"
bounds_path.write_text(json.dumps(bounds, separators=(",", ":")), encoding="utf-8")
elapsed_ms = (time.perf_counter() - started) * 1000.0
message = f"Image contour finished in {elapsed_ms:.0f} ms"
message = f"Image background removal finished in {elapsed_ms:.0f} ms"
context.progress(1.0, message)
return {"image": output_path.name, "bounds": bounds_path.name}