limbicnation
faab680db5
- Remove unused imports (ruff F401) - Fix f-strings without placeholders (ruff F541) - Resolve E402 import ordering in grabcut_nodes.py - Fix B108 hardcoded /tmp fallback in nodes.py (use tempfile.gettempdir()) - Delete untracked artifacts: .claude/, node.zip, __pycache__/ - Add .claude/ and *.zip to .gitignore
412 lines
14 KiB
Python
412 lines
14 KiB
Python
"""
|
|
RemoveBackgroundAndResizeNode
|
|
============================
|
|
Single-file ComfyUI node for background removal with integrated resize.
|
|
|
|
Based on ``TransparencyBackgroundRemover`` and ``ScalingMixin`` from the
|
|
ComfyUI-TransparencyBackgroundRemover package. Adapted for standalone use
|
|
with minimal dependencies — PIL only for background removal, no hard
|
|
OpenCV requirement.
|
|
|
|
Author: Gero Doll / Limbicnation
|
|
License: Apache-2.0
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
try:
|
|
import structlog
|
|
log = structlog.get_logger(__name__)
|
|
except ImportError:
|
|
log = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ScalingMixin — shared resize logic (see scaling.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
try:
|
|
from .scaling import ScalingMixin # noqa: F401 — re-exported for convenience
|
|
except ImportError:
|
|
from scaling import ScalingMixin
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PIL-only background removal
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _kmeans_colors(
|
|
arr: np.ndarray, n_clusters: int, tol: float = 0.01, max_iter: int = 20,
|
|
max_samples: int = 50000,
|
|
) -> np.ndarray:
|
|
"""
|
|
Simple k-means on pixel array [H*W, channels].
|
|
Returns ``(n_clusters, channels)`` cluster centres.
|
|
Fallback when scikit-learn is unavailable.
|
|
|
|
Parameters
|
|
----------
|
|
arr : np.ndarray
|
|
Input image array (H, W, C).
|
|
n_clusters : int
|
|
Number of k-means clusters.
|
|
tol : float
|
|
Convergence tolerance for centre shift.
|
|
max_iter : int
|
|
Maximum k-means iterations.
|
|
max_samples : int
|
|
Cap on number of pixels fed to k-means to avoid OOM on large images.
|
|
Subsampled uniformly when pixel count exceeds this value.
|
|
"""
|
|
pixels = arr.reshape(-1, arr.shape[-1]).astype(np.float32)
|
|
|
|
# Subsample to avoid OOM on large images (e.g. 4K → 8.3M pixels)
|
|
if len(pixels) > max_samples:
|
|
rng = np.random.RandomState(42)
|
|
idx = rng.choice(len(pixels), size=max_samples, replace=False)
|
|
pixels = pixels[idx]
|
|
|
|
indices = np.random.RandomState(42).choice(
|
|
len(pixels), size=n_clusters, replace=False
|
|
)
|
|
centres = pixels[indices].copy()
|
|
|
|
for _ in range(max_iter):
|
|
dists = np.linalg.norm(pixels[:, None] - centres[None], axis=-1)
|
|
labels = dists.argmin(axis=1)
|
|
new_centres = np.stack([pixels[labels == i].mean(axis=0) for i in range(n_clusters)])
|
|
if np.abs(new_centres - centres).max() < tol:
|
|
break
|
|
centres = new_centres
|
|
|
|
return centres
|
|
|
|
|
|
def _dominant_bg(arr: np.ndarray, samples: int = 5000) -> np.ndarray:
|
|
"""Sample corners + border to estimate background colour."""
|
|
h, w = arr.shape[:2]
|
|
mask = np.zeros((h, w), dtype=bool)
|
|
bw = max(1, min(h, w) // 10)
|
|
|
|
# top / bottom strips and left / right strips (intersection gives the four corners)
|
|
mask[:bw, :] = True
|
|
mask[-bw:, :] = True
|
|
mask[:, :bw] = True
|
|
mask[:, -bw:] = True
|
|
|
|
bg_pixels = arr[mask]
|
|
if len(bg_pixels) > samples:
|
|
idx = np.random.RandomState(42).choice(len(bg_pixels), samples, replace=False)
|
|
bg_pixels = bg_pixels[idx]
|
|
|
|
return bg_pixels.mean(axis=0).astype(np.uint8)
|
|
|
|
|
|
def remove_background_pil(
|
|
image: Image.Image,
|
|
tolerance: int = 30,
|
|
color_clusters: int = 8,
|
|
foreground_bias: float = 0.7,
|
|
binary_threshold: int = 128,
|
|
dither_handling: bool = True,
|
|
) -> Image.Image:
|
|
"""
|
|
Remove background from a PIL Image using colour clustering.
|
|
|
|
Pipeline
|
|
--------
|
|
1. Flatten to ``[H*W, channels]`` and run k-means to find dominant colours.
|
|
2. Identify background as the colour most similar to image edges / corners.
|
|
3. Compute per-pixel distance to foreground centres; keep pixels that are
|
|
closer than ``tolerance`` (adjusted by ``foreground_bias``).
|
|
4. Apply simple edge refinement via binary thresholding.
|
|
5. Return ``RGBA`` PIL Image.
|
|
|
|
Parameters
|
|
----------
|
|
image : PIL.Image
|
|
Input image (RGB or RGBA; RGBA ``A`` is ignored).
|
|
tolerance : int
|
|
Colour distance threshold (0-255). Higher = more pixels accepted as foreground.
|
|
color_clusters : int
|
|
K-means centres used to identify distinct colours.
|
|
foreground_bias : float
|
|
Bias towards keeping pixels (0 = discard more, 1 = keep more).
|
|
binary_threshold : int
|
|
Threshold for alpha mask binarisation (0-255).
|
|
dither_handling : bool
|
|
When True, apply a 1-pixel border erosion to reduce dither artefacts.
|
|
|
|
Returns
|
|
-------
|
|
PIL.Image
|
|
RGBA image with transparent background.
|
|
"""
|
|
rgb = image.convert("RGB")
|
|
arr = np.array(rgb, dtype=np.float32) # [H, W, 3]
|
|
|
|
# 1. k-means colour clustering
|
|
if color_clusters < 2:
|
|
color_clusters = 2
|
|
centres = _kmeans_colors(arr, color_clusters)
|
|
|
|
# 2. background colour from image perimeter
|
|
bg_colour = _dominant_bg(arr)
|
|
|
|
# 3. distance of each pixel to background vs foreground centres
|
|
h, w = arr.shape[:2]
|
|
flat = arr.reshape(-1, 3)
|
|
|
|
bg_dist = np.linalg.norm(flat - bg_colour.astype(np.float32), axis=1)
|
|
fg_dist = np.linalg.norm(
|
|
flat[:, None] - centres.astype(np.float32), axis=-1
|
|
).min(axis=1)
|
|
|
|
# effective tolerance — foreground_bias shifts the acceptance boundary
|
|
effective_tol = tolerance * (1.5 - foreground_bias)
|
|
mask = (fg_dist + effective_tol) < bg_dist
|
|
|
|
alpha = mask.reshape(h, w).astype(np.uint8) * 255
|
|
|
|
# 4. edge refinement
|
|
if dither_handling:
|
|
from PIL import ImageFilter
|
|
|
|
alpha_pil = Image.fromarray(alpha, mode="L")
|
|
alpha_pil = alpha_pil.filter(ImageFilter.MinFilter(3))
|
|
alpha = np.array(alpha_pil)
|
|
|
|
# 5. binarise
|
|
alpha = (alpha > binary_threshold).astype(np.uint8) * 255
|
|
|
|
# 6. assemble RGBA
|
|
result = np.dstack([arr.astype(np.uint8), alpha])
|
|
return Image.fromarray(result, mode="RGBA")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ComfyUI node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class RemoveBackgroundAndResizeNode(ScalingMixin):
|
|
"""
|
|
Combines background removal and resize in a single node.
|
|
|
|
``INPUT_TYPES`` exposes tolerance, edge sensitivity, colour-cluster, and
|
|
output-size controls so the node can be dropped into any workflow without
|
|
separate background-removal and resize chains.
|
|
|
|
Outputs
|
|
-------
|
|
IMAGE : RGBA tensor (N, H, W, 4), values 0-1
|
|
MASK : greyscale alpha tensor (N, H, W), values 0-1
|
|
"""
|
|
|
|
# ------------------------------------------------------------------
|
|
# ComfyUI contract
|
|
# ------------------------------------------------------------------
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image/processing"
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {"tooltip": "Input image tensor from upstream node"}),
|
|
"tolerance": (
|
|
"INT",
|
|
{
|
|
"default": 30,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": (
|
|
"Colour distance threshold. Higher = more aggressive "
|
|
"foreground extraction (0-255)."
|
|
),
|
|
},
|
|
),
|
|
"foreground_bias": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.7,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "slider",
|
|
"tooltip": (
|
|
"Bias toward keeping pixels as foreground. "
|
|
"Higher = more pixels preserved."
|
|
),
|
|
},
|
|
),
|
|
"color_clusters": (
|
|
"INT",
|
|
{
|
|
"default": 8,
|
|
"min": 2,
|
|
"max": 20,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Number of k-means centres for colour segmentation.",
|
|
},
|
|
),
|
|
"output_size": (
|
|
[
|
|
"ORIGINAL",
|
|
"512x512",
|
|
"768x768",
|
|
"1024x1024",
|
|
"1280x720",
|
|
"1920x1080",
|
|
],
|
|
{
|
|
"default": "ORIGINAL",
|
|
"tooltip": (
|
|
"Resize output to this size. "
|
|
"ORIGINAL keeps source dimensions."
|
|
),
|
|
},
|
|
),
|
|
"scaling_method": (
|
|
["NEAREST", "BILINEAR", "BICUBIC", "LANCZOS"],
|
|
{
|
|
"default": "LANCZOS",
|
|
"tooltip": (
|
|
"Interpolation method. "
|
|
"LANCZOS = best quality; "
|
|
"NEAREST = pixel-perfect for pixel art."
|
|
),
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"dither_handling": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Apply edge erosion to suppress dither artefacts.",
|
|
},
|
|
),
|
|
"binary_threshold": (
|
|
"INT",
|
|
{
|
|
"default": 128,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Alpha mask binarisation cutoff (0-255).",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# processing
|
|
# ------------------------------------------------------------------
|
|
|
|
@torch.no_grad()
|
|
def process(
|
|
self,
|
|
image: torch.Tensor,
|
|
tolerance: int,
|
|
foreground_bias: float,
|
|
color_clusters: int,
|
|
output_size: str,
|
|
scaling_method: str,
|
|
dither_handling: bool = True,
|
|
binary_threshold: int = 128,
|
|
) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
|
|
"""
|
|
``image`` — ComfyUI IMAGE tensor, shape (N, H, W, C), float32 0-1.
|
|
Returns (IMAGE tensor, MASK tensor).
|
|
"""
|
|
# ------------------------------------------------------------------
|
|
# validation
|
|
# ------------------------------------------------------------------
|
|
if image is None or image.numel() == 0:
|
|
raise ValueError("RemoveBackgroundAndResizeNode: received empty image tensor.")
|
|
if image.ndim != 4:
|
|
raise ValueError(
|
|
f"RemoveBackgroundAndResizeNode: expected 4-D IMAGE tensor "
|
|
f"(N,H,W,C), got {image.ndim}-D."
|
|
)
|
|
if image.shape[-1] not in (3, 4):
|
|
raise ValueError(
|
|
f"RemoveBackgroundAndResizeNode: IMAGE must be RGB or RGBA; "
|
|
f"got channels={image.shape[-1]}."
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# per-batch processing
|
|
# ------------------------------------------------------------------
|
|
rgba_tensors: list[torch.Tensor] = []
|
|
mask_tensors: list[torch.Tensor] = []
|
|
|
|
for i in range(image.shape[0]):
|
|
img_np = (image[i].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
|
|
|
# handle RGB input
|
|
if img_np.shape[-1] == 3:
|
|
pil_in = Image.fromarray(img_np, mode="RGB")
|
|
else:
|
|
pil_in = Image.fromarray(img_np[:, :, :3], mode="RGB")
|
|
|
|
# ---- remove background ----
|
|
pil_rgba = remove_background_pil(
|
|
pil_in,
|
|
tolerance=tolerance,
|
|
foreground_bias=foreground_bias,
|
|
color_clusters=color_clusters,
|
|
binary_threshold=binary_threshold,
|
|
dither_handling=dither_handling,
|
|
)
|
|
|
|
# ---- resize ----
|
|
pil_rgba = self.resize_image(pil_rgba, output_size, scaling_method)
|
|
|
|
# ---- back to tensor ----
|
|
arr = np.array(pil_rgba, dtype=np.float32) / 255.0
|
|
rgba_tensors.append(torch.from_numpy(arr))
|
|
|
|
# extract alpha channel as mask
|
|
mask_np = (np.array(pil_rgba.split()[3], dtype=np.float32) / 255.0)
|
|
mask_tensors.append(torch.from_numpy(mask_np))
|
|
|
|
# ------------------------------------------------------------------
|
|
# stack and return in ComfyUI tensor format
|
|
# ------------------------------------------------------------------
|
|
stacked_image: torch.Tensor = torch.stack(rgba_tensors).to(dtype=torch.float32)
|
|
stacked_mask: torch.Tensor = torch.stack(mask_tensors).to(dtype=torch.float32)
|
|
|
|
# ComfyUI MASK convention: (N, H, W)
|
|
if stacked_mask.ndim == 4 and stacked_mask.shape[-1] == 1:
|
|
stacked_mask = stacked_mask.squeeze(-1)
|
|
|
|
log.debug(
|
|
"RemoveBackgroundAndResizeNode.done",
|
|
batch=image.shape[0],
|
|
output_size=output_size,
|
|
scaling_method=scaling_method,
|
|
)
|
|
|
|
return (stacked_image, stacked_mask)
|
|
|
|
|
|
# ComfyUI registration dict
|
|
NODE_CLASS_MAPPINGS = {"RemoveBackgroundAndResize": RemoveBackgroundAndResizeNode}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RemoveBackgroundAndResize": "Remove Background & Resize"
|
|
}
|