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