217 lines
6.4 KiB
Python
217 lines
6.4 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
|
|
try:
|
|
import numpy as np
|
|
except Exception as exc: # pragma: no cover - optional runtime dependency
|
|
np = None
|
|
_NUMPY_ERR = exc
|
|
else: # pragma: no cover - exercised only in ComfyUI-like runtime
|
|
_NUMPY_ERR = None
|
|
|
|
try:
|
|
from PIL import Image
|
|
except Exception as exc: # pragma: no cover - optional runtime dependency
|
|
Image = None
|
|
_PIL_ERR = exc
|
|
else: # pragma: no cover - exercised only in ComfyUI-like runtime
|
|
_PIL_ERR = None
|
|
|
|
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
|
|
|
|
|
|
_DEFAULT_MASK_STATE = {
|
|
"version": 1,
|
|
"mask_png_base64": "",
|
|
"touched": False,
|
|
"editor": {},
|
|
}
|
|
|
|
|
|
def parse_mask_state(mask_state_json):
|
|
if not mask_state_json:
|
|
return dict(_DEFAULT_MASK_STATE)
|
|
|
|
try:
|
|
payload = json.loads(mask_state_json)
|
|
except (TypeError, ValueError):
|
|
return dict(_DEFAULT_MASK_STATE)
|
|
|
|
if not isinstance(payload, dict):
|
|
return dict(_DEFAULT_MASK_STATE)
|
|
|
|
state = dict(_DEFAULT_MASK_STATE)
|
|
state["version"] = _coerce_int(payload.get("version"), 1)
|
|
state["mask_png_base64"] = str(payload.get("mask_png_base64") or "")
|
|
state["touched"] = _coerce_bool(payload.get("touched"), False)
|
|
state["editor"] = _coerce_dict(payload.get("editor"))
|
|
return state
|
|
|
|
|
|
def resolve_output_mask(image, input_mask, mask_state_json, invert_mask):
|
|
state = parse_mask_state(mask_state_json)
|
|
resolved = None
|
|
|
|
if state["touched"]:
|
|
resolved = _decode_mask_png_base64(state["mask_png_base64"], image)
|
|
|
|
if resolved is None and input_mask is not None:
|
|
resolved = _resize_mask_to_image(input_mask, image)
|
|
|
|
if resolved is None:
|
|
resolved = _make_black_mask(image)
|
|
|
|
if invert_mask:
|
|
resolved = 1.0 - resolved
|
|
|
|
return resolved.clamp(0.0, 1.0)
|
|
|
|
|
|
def build_preview_image(image, mask, overlay_alpha):
|
|
_require_torch("[LLS] torch is required for preview image generation.")
|
|
|
|
alpha = max(0.0, min(1.0, _coerce_float(overlay_alpha, 0.4)))
|
|
normalized_mask = _resize_mask_to_image(mask, image).unsqueeze(-1)
|
|
overlay_strength = normalized_mask * alpha
|
|
overlay = torch.zeros_like(image)
|
|
overlay[..., 0] = 1.0
|
|
return (image * (1.0 - overlay_strength) + (overlay * overlay_strength)).clamp(0.0, 1.0)
|
|
|
|
|
|
def _get_image_size(image):
|
|
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 = _coerce_int(width, 0)
|
|
height = _coerce_int(height, 0)
|
|
if width <= 0 or height <= 0:
|
|
raise RuntimeError("[LLS] image width and height must be positive.")
|
|
return width, height
|
|
|
|
|
|
def _get_image_batch_size(image):
|
|
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].")
|
|
|
|
batch = _coerce_int(shape[0], 0)
|
|
if batch <= 0:
|
|
raise RuntimeError("[LLS] image batch size must be positive.")
|
|
return batch
|
|
|
|
|
|
def _make_black_mask(image):
|
|
_require_torch("[LLS] torch is required for mask generation.")
|
|
|
|
width, height = _get_image_size(image)
|
|
return torch.zeros((image.shape[0], height, width), dtype=image.dtype, device=image.device)
|
|
|
|
|
|
def _resize_mask_to_image(mask, image):
|
|
_require_torch("[LLS] torch is required for mask resizing.")
|
|
|
|
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 = _get_image_batch_size(image)
|
|
mask_batch = _coerce_int(shape[0], 0)
|
|
if mask_batch <= 0:
|
|
raise RuntimeError("[LLS] mask batch size must be positive.")
|
|
|
|
width, height = _get_image_size(image)
|
|
normalized = mask.to(device=image.device, dtype=image.dtype).clamp(0.0, 1.0)
|
|
if mask_batch != image_batch:
|
|
if mask_batch == 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 {mask_batch} and image batch {image_batch}."
|
|
)
|
|
if tuple(normalized.shape[1:]) == (height, width):
|
|
return normalized
|
|
|
|
resized = torch_nn_functional.interpolate(
|
|
normalized.unsqueeze(1),
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
return resized.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
|
|
def _decode_mask_png_base64(mask_png_base64, image):
|
|
_require_torch("[LLS] torch is required for mask decoding.")
|
|
_require_mask_decode_deps()
|
|
|
|
if not mask_png_base64:
|
|
return None
|
|
|
|
width, height = _get_image_size(image)
|
|
try:
|
|
payload = base64.b64decode(mask_png_base64)
|
|
decoded = Image.open(io.BytesIO(payload)).convert("L")
|
|
except Exception:
|
|
return None
|
|
|
|
if decoded.size != (width, height):
|
|
decoded = decoded.resize((width, height), Image.BILINEAR)
|
|
|
|
array = np.asarray(decoded, dtype=np.float32) / 255.0
|
|
tensor = torch.from_numpy(array).to(device=image.device, dtype=image.dtype)
|
|
tensor = tensor.unsqueeze(0)
|
|
if image.shape[0] > 1:
|
|
tensor = tensor.repeat(image.shape[0], 1, 1)
|
|
return tensor.clamp(0.0, 1.0)
|
|
|
|
|
|
def _require_torch(message):
|
|
if torch is None or torch_nn_functional is None:
|
|
raise RuntimeError(message) from _TORCH_ERR
|
|
|
|
|
|
def _require_mask_decode_deps():
|
|
if np is None:
|
|
raise RuntimeError("[LLS] numpy is required for mask decoding.") from _NUMPY_ERR
|
|
if Image is None:
|
|
raise RuntimeError("[LLS] Pillow is required for mask decoding.") from _PIL_ERR
|
|
|
|
|
|
def _coerce_bool(value, default):
|
|
if isinstance(value, bool):
|
|
return value
|
|
return default
|
|
|
|
|
|
def _coerce_dict(value):
|
|
if isinstance(value, dict):
|
|
return dict(value)
|
|
return {}
|
|
|
|
|
|
def _coerce_int(value, default):
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _coerce_float(value, default):
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return default
|