333 lines
11 KiB
Python
333 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
try:
|
|
import torch
|
|
import torch.nn.functional as torch_nn_functional
|
|
except Exception as exc: # pragma: no cover - optional runtime dependency
|
|
torch = None
|
|
torch_nn_functional = None
|
|
_TORCH_ERR = exc
|
|
else: # pragma: no cover - exercised only in ComfyUI-like runtime
|
|
_TORCH_ERR = None
|
|
|
|
|
|
MASK_INFO_TYPE = "LLS_MASK_INFO"
|
|
SHAPE_TYPE_CHOICES = ["rectangle", "square", "circle", "ellipse"]
|
|
COORDINATE_MODE_CHOICES = ["pixel", "percent"]
|
|
COMBINE_MODE_CHOICES = ["replace", "union", "subtract", "intersect"]
|
|
OVERLAY_COLOR_CHOICES = ["red", "green", "blue"]
|
|
|
|
_OVERLAY_COLORS = {
|
|
"red": (1.0, 0.0, 0.0),
|
|
"green": (0.0, 1.0, 0.0),
|
|
"blue": (0.0, 0.0, 1.0),
|
|
}
|
|
|
|
|
|
def create_reference_image(
|
|
image_width: int,
|
|
image_height: int,
|
|
*,
|
|
batch: int = 1,
|
|
dtype=None,
|
|
device=None,
|
|
):
|
|
_require_torch("[LLS] torch is required for mask generation.")
|
|
|
|
width = max(1, int(image_width))
|
|
height = max(1, int(image_height))
|
|
batch = max(1, int(batch))
|
|
kwargs = {"dtype": dtype or torch.float32}
|
|
if device is not None:
|
|
kwargs["device"] = device
|
|
return torch.zeros((batch, height, width, 3), **kwargs)
|
|
|
|
|
|
def mask_to_image(mask):
|
|
_require_torch("[LLS] torch is required for mask image generation.")
|
|
|
|
shape = getattr(mask, "shape", None)
|
|
if not isinstance(shape, (list, tuple)) or len(shape) != 3:
|
|
raise RuntimeError("[LLS] mask must have shape [batch, height, width].")
|
|
normalized = mask.clamp(0.0, 1.0)
|
|
return normalized.unsqueeze(-1).repeat(1, 1, 1, 3)
|
|
|
|
|
|
def get_image_size(image) -> tuple[int, int]:
|
|
shape = getattr(image, "shape", None)
|
|
if not isinstance(shape, (list, tuple)) or len(shape) != 4:
|
|
raise RuntimeError("[LLS] image must have shape [batch, height, width, channels].")
|
|
_, height, width, _ = tuple(shape)
|
|
width = int(width)
|
|
height = int(height)
|
|
if width <= 0 or height <= 0:
|
|
raise RuntimeError("[LLS] image width and height must be positive.")
|
|
return width, height
|
|
|
|
|
|
def create_shape_mask(
|
|
image,
|
|
*,
|
|
shape_type: str,
|
|
coordinate_mode: str,
|
|
center_x: float,
|
|
center_y: float,
|
|
width: float,
|
|
height: float,
|
|
radius: float,
|
|
):
|
|
_require_torch("[LLS] torch is required for mask generation.")
|
|
|
|
image_width, image_height = get_image_size(image)
|
|
batch = int(image.shape[0])
|
|
dtype = image.dtype
|
|
device = image.device
|
|
|
|
shape_type = str(shape_type or "rectangle")
|
|
coordinate_mode = str(coordinate_mode or "percent")
|
|
center_px_x = _resolve_axis_value(center_x, image_width, coordinate_mode)
|
|
center_px_y = _resolve_axis_value(center_y, image_height, coordinate_mode)
|
|
width_px = _resolve_length_value(width, image_width, coordinate_mode)
|
|
height_px = _resolve_length_value(height, image_height, coordinate_mode)
|
|
radius_px = _resolve_radius_value(radius, min(image_width, image_height), coordinate_mode)
|
|
|
|
if shape_type == "square":
|
|
width_px = max(1.0, width_px)
|
|
height_px = width_px
|
|
if shape_type == "circle":
|
|
radius_px = max(1.0, radius_px)
|
|
if shape_type not in {"rectangle", "square", "circle", "ellipse"}:
|
|
raise RuntimeError(f"[LLS] Unsupported shape_type '{shape_type}'.")
|
|
|
|
xs = torch.arange(image_width, dtype=dtype, device=device) + 0.5
|
|
ys = torch.arange(image_height, dtype=dtype, device=device) + 0.5
|
|
try:
|
|
yy, xx = torch.meshgrid(ys, xs, indexing="ij")
|
|
except TypeError: # pragma: no cover - older torch fallback
|
|
yy, xx = torch.meshgrid(ys, xs)
|
|
|
|
if shape_type in {"rectangle", "square"}:
|
|
half_width = width_px / 2.0
|
|
half_height = height_px / 2.0
|
|
x1 = center_px_x - half_width
|
|
y1 = center_px_y - half_height
|
|
x2 = center_px_x + half_width
|
|
y2 = center_px_y + half_height
|
|
inside = (xx >= x1) & (xx < x2) & (yy >= y1) & (yy < y2)
|
|
geometric_area_px = float(width_px * height_px)
|
|
unclipped_bbox = (x1, y1, x2, y2)
|
|
elif shape_type == "circle":
|
|
inside = ((xx - center_px_x) ** 2) + ((yy - center_px_y) ** 2) <= (radius_px ** 2)
|
|
geometric_area_px = float(math.pi * radius_px * radius_px)
|
|
unclipped_bbox = (
|
|
center_px_x - radius_px,
|
|
center_px_y - radius_px,
|
|
center_px_x + radius_px,
|
|
center_px_y + radius_px,
|
|
)
|
|
else:
|
|
rx = max(0.5, width_px / 2.0)
|
|
ry = max(0.5, height_px / 2.0)
|
|
inside = (((xx - center_px_x) / rx) ** 2) + (((yy - center_px_y) / ry) ** 2) <= 1.0
|
|
geometric_area_px = float(math.pi * rx * ry)
|
|
unclipped_bbox = (
|
|
center_px_x - rx,
|
|
center_px_y - ry,
|
|
center_px_x + rx,
|
|
center_px_y + ry,
|
|
)
|
|
|
|
mask_2d = inside.to(dtype=dtype)
|
|
mask = mask_2d.unsqueeze(0).repeat(batch, 1, 1)
|
|
clipped_bbox, clipped_by_image = _clip_bbox(unclipped_bbox, image_width, image_height)
|
|
|
|
geometry_info = {
|
|
"image_size": [image_width, image_height],
|
|
"shape_type": shape_type,
|
|
"coordinate_mode": coordinate_mode,
|
|
"center": [float(center_px_x), float(center_px_y)],
|
|
"width": float(width_px) if shape_type in {"rectangle", "square", "ellipse"} else None,
|
|
"height": float(height_px) if shape_type in {"rectangle", "square", "ellipse"} else None,
|
|
"radius": float(radius_px) if shape_type == "circle" else None,
|
|
"bbox": clipped_bbox,
|
|
"geometric_area_px": float(geometric_area_px),
|
|
"clipped_by_image": bool(clipped_by_image),
|
|
}
|
|
return mask, geometry_info
|
|
|
|
|
|
def align_mask_to_image(mask, image):
|
|
_require_torch("[LLS] torch is required for mask resizing.")
|
|
|
|
if mask is None:
|
|
return None
|
|
|
|
shape = getattr(mask, "shape", None)
|
|
if not isinstance(shape, (list, tuple)) or len(shape) != 3:
|
|
raise RuntimeError("[LLS] mask must have shape [batch, height, width].")
|
|
|
|
image_batch = int(image.shape[0])
|
|
image_width, image_height = get_image_size(image)
|
|
normalized = mask.to(device=image.device, dtype=image.dtype).clamp(0.0, 1.0)
|
|
if int(normalized.shape[0]) != image_batch:
|
|
if int(normalized.shape[0]) == 1 and image_batch > 1:
|
|
normalized = normalized.repeat(image_batch, 1, 1)
|
|
else:
|
|
raise RuntimeError(
|
|
f"[LLS] mask batch size must match image batch size or be 1; got mask batch {int(normalized.shape[0])} and image batch {image_batch}."
|
|
)
|
|
|
|
if tuple(normalized.shape[1:]) == (image_height, image_width):
|
|
return normalized
|
|
|
|
resized = torch_nn_functional.interpolate(
|
|
normalized.unsqueeze(1),
|
|
size=(image_height, image_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
return resized.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
|
|
def combine_masks(created_mask, input_mask, image, combine_mode: str):
|
|
aligned_input = align_mask_to_image(input_mask, image) if input_mask is not None else None
|
|
if aligned_input is None:
|
|
return created_mask
|
|
|
|
combine_mode = str(combine_mode or "replace")
|
|
if combine_mode == "replace":
|
|
return created_mask
|
|
if combine_mode == "union":
|
|
return torch.maximum(aligned_input, created_mask)
|
|
if combine_mode == "subtract":
|
|
return torch.clamp(aligned_input - created_mask, 0.0, 1.0)
|
|
if combine_mode == "intersect":
|
|
return aligned_input * created_mask
|
|
raise RuntimeError(f"[LLS] Unsupported combine_mode '{combine_mode}'.")
|
|
|
|
|
|
def apply_invert(mask, invert_mask: bool):
|
|
if not invert_mask:
|
|
return mask
|
|
return 1.0 - mask
|
|
|
|
|
|
def apply_softening(mask, *, feather: float, blur: float):
|
|
feather_radius = _coerce_radius(feather, limit=128.0)
|
|
blur_radius = _coerce_radius(blur, limit=128.0)
|
|
softened = mask
|
|
if feather_radius > 0.0:
|
|
softened = _box_blur_mask(softened, feather_radius)
|
|
if blur_radius > 0.0:
|
|
softened = _box_blur_mask(softened, blur_radius)
|
|
return softened.clamp(0.0, 1.0)
|
|
|
|
|
|
def build_preview_image(image, mask, *, overlay_alpha: float, overlay_color: str):
|
|
_require_torch("[LLS] torch is required for preview image generation.")
|
|
|
|
overlay_alpha = max(0.0, min(1.0, float(overlay_alpha)))
|
|
overlay_rgb = _OVERLAY_COLORS.get(str(overlay_color or "red"), _OVERLAY_COLORS["red"])
|
|
overlay = torch.zeros_like(image)
|
|
overlay[..., 0] = overlay_rgb[0]
|
|
overlay[..., 1] = overlay_rgb[1]
|
|
overlay[..., 2] = overlay_rgb[2]
|
|
|
|
mask_rgb = align_mask_to_image(mask, image).unsqueeze(-1)
|
|
strength = mask_rgb * overlay_alpha
|
|
return (image * (1.0 - strength) + overlay * strength).clamp(0.0, 1.0)
|
|
|
|
|
|
def build_area_info(
|
|
mask,
|
|
geometry_info: dict,
|
|
*,
|
|
feather: float,
|
|
blur: float,
|
|
invert_mask: bool,
|
|
combine_mode: str,
|
|
):
|
|
_require_torch("[LLS] torch is required for area_info calculation.")
|
|
|
|
image_width, image_height = geometry_info["image_size"]
|
|
binary_area_px = int((mask > 0.5).sum().item())
|
|
effective_area_px = float(mask.sum().item())
|
|
total_area = max(1, image_width * image_height)
|
|
|
|
info = dict(geometry_info)
|
|
info.update(
|
|
{
|
|
"binary_area_px": binary_area_px,
|
|
"effective_area_px": effective_area_px,
|
|
"area_ratio": float(binary_area_px / float(total_area)),
|
|
"feather": float(_coerce_radius(feather, limit=128.0)),
|
|
"blur": float(_coerce_radius(blur, limit=128.0)),
|
|
"invert_mask": bool(invert_mask),
|
|
"combine_mode": str(combine_mode or "replace"),
|
|
}
|
|
)
|
|
return info
|
|
|
|
|
|
def _box_blur_mask(mask, radius: float):
|
|
kernel_radius = max(1, int(round(float(radius))))
|
|
kernel_size = (kernel_radius * 2) + 1
|
|
blurred = torch_nn_functional.avg_pool2d(
|
|
mask.unsqueeze(1),
|
|
kernel_size=kernel_size,
|
|
stride=1,
|
|
padding=kernel_radius,
|
|
)
|
|
return blurred.squeeze(1)
|
|
|
|
|
|
def _clip_bbox(unclipped_bbox, image_width: int, image_height: int):
|
|
x1, y1, x2, y2 = [float(value) for value in unclipped_bbox]
|
|
clipped_x1 = max(0, min(int(math.floor(x1)), image_width))
|
|
clipped_y1 = max(0, min(int(math.floor(y1)), image_height))
|
|
clipped_x2 = max(clipped_x1, min(int(math.ceil(x2)), image_width))
|
|
clipped_y2 = max(clipped_y1, min(int(math.ceil(y2)), image_height))
|
|
clipped_by_image = (
|
|
clipped_x1 != int(math.floor(x1))
|
|
or clipped_y1 != int(math.floor(y1))
|
|
or clipped_x2 != int(math.ceil(x2))
|
|
or clipped_y2 != int(math.ceil(y2))
|
|
)
|
|
return [clipped_x1, clipped_y1, clipped_x2, clipped_y2], clipped_by_image
|
|
|
|
|
|
def _resolve_axis_value(value: float, image_extent: int, coordinate_mode: str) -> float:
|
|
value = float(value)
|
|
if coordinate_mode == "percent":
|
|
return value * float(image_extent)
|
|
return value
|
|
|
|
|
|
def _resolve_length_value(value: float, image_extent: int, coordinate_mode: str) -> float:
|
|
value = float(value)
|
|
if coordinate_mode == "percent":
|
|
return max(1.0, value * float(image_extent))
|
|
return max(1.0, value)
|
|
|
|
|
|
def _resolve_radius_value(value: float, image_extent: int, coordinate_mode: str) -> float:
|
|
value = float(value)
|
|
if coordinate_mode == "percent":
|
|
return max(1.0, value * float(image_extent))
|
|
return max(1.0, value)
|
|
|
|
|
|
def _coerce_radius(value: float, *, limit: float) -> float:
|
|
try:
|
|
normalized = float(value)
|
|
except (TypeError, ValueError):
|
|
return 0.0
|
|
return max(0.0, min(limit, normalized))
|
|
|
|
|
|
def _require_torch(message: str):
|
|
if torch is None or torch_nn_functional is None:
|
|
raise RuntimeError(message) from _TORCH_ERR
|