254 lines
7.3 KiB
Python
254 lines
7.3 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
import math
|
|
|
|
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 in runtime/tests
|
|
_NUMPY_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 in runtime/tests
|
|
_TORCH_ERR = None
|
|
|
|
|
|
def _require_torch():
|
|
if torch is None or torch_nn_functional is None:
|
|
raise RuntimeError("[LLS] torch is required for mask processing nodes.") from _TORCH_ERR
|
|
|
|
|
|
def _coerce_non_negative_int(value) -> int:
|
|
try:
|
|
return max(0, int(value))
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def normalize_mask(mask):
|
|
_require_torch()
|
|
|
|
if mask is None:
|
|
raise RuntimeError("[LLS] mask input is required.")
|
|
if not isinstance(mask, torch.Tensor):
|
|
mask = torch.as_tensor(mask)
|
|
|
|
ndim = int(mask.ndim)
|
|
if ndim == 2:
|
|
normalized = mask.unsqueeze(0)
|
|
elif ndim == 3:
|
|
normalized = mask
|
|
elif ndim == 4:
|
|
if int(mask.shape[-1]) == 1:
|
|
normalized = mask.squeeze(-1)
|
|
elif int(mask.shape[1]) == 1:
|
|
normalized = mask.squeeze(1)
|
|
else:
|
|
raise RuntimeError(
|
|
"[LLS] mask must have shape [H,W], [B,H,W], [B,H,W,1], or [B,1,H,W]."
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
"[LLS] mask must have shape [H,W], [B,H,W], [B,H,W,1], or [B,1,H,W]."
|
|
)
|
|
|
|
return normalized.to(dtype=torch.float32).clamp(0.0, 1.0)
|
|
|
|
|
|
def resize_mask_to_hw(mask, height, width):
|
|
normalized = normalize_mask(mask)
|
|
target_height = max(1, int(height))
|
|
target_width = max(1, int(width))
|
|
if tuple(normalized.shape[1:]) == (target_height, target_width):
|
|
return normalized
|
|
|
|
padded = normalized.unsqueeze(1)
|
|
resized = torch_nn_functional.interpolate(
|
|
padded,
|
|
size=(target_height, target_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
return resized.squeeze(1).clamp(0.0, 1.0).to(dtype=torch.float32)
|
|
|
|
|
|
def align_mask_batch(mask, target_batch):
|
|
normalized = normalize_mask(mask)
|
|
batch = int(normalized.shape[0])
|
|
target = max(1, int(target_batch))
|
|
if batch == target:
|
|
return normalized
|
|
if batch == 1:
|
|
return normalized.repeat(target, 1, 1)
|
|
if batch < target:
|
|
repeats = int(math.ceil(target / batch))
|
|
return normalized.repeat(repeats, 1, 1)[:target]
|
|
return normalized[:target]
|
|
|
|
|
|
def apply_grow(mask, radius):
|
|
normalized = normalize_mask(mask)
|
|
grow_radius = _coerce_non_negative_int(radius)
|
|
if grow_radius == 0:
|
|
return normalized
|
|
|
|
kernel_size = (grow_radius * 2) + 1
|
|
grown = torch_nn_functional.max_pool2d(
|
|
normalized.unsqueeze(1),
|
|
kernel_size=kernel_size,
|
|
stride=1,
|
|
padding=grow_radius,
|
|
)
|
|
return grown.squeeze(1).clamp(0.0, 1.0).to(dtype=torch.float32)
|
|
|
|
|
|
def apply_shrink(mask, radius):
|
|
normalized = normalize_mask(mask)
|
|
shrink_radius = _coerce_non_negative_int(radius)
|
|
if shrink_radius == 0:
|
|
return normalized
|
|
|
|
inverted = 1.0 - normalized
|
|
kernel_size = (shrink_radius * 2) + 1
|
|
pooled = torch_nn_functional.max_pool2d(
|
|
inverted.unsqueeze(1),
|
|
kernel_size=kernel_size,
|
|
stride=1,
|
|
padding=shrink_radius,
|
|
)
|
|
return (1.0 - pooled.squeeze(1)).clamp(0.0, 1.0).to(dtype=torch.float32)
|
|
|
|
|
|
def apply_blur(mask, radius):
|
|
normalized = normalize_mask(mask)
|
|
blur_radius = _coerce_non_negative_int(radius)
|
|
if blur_radius == 0:
|
|
return normalized
|
|
|
|
kernel_size = (blur_radius * 2) + 1
|
|
padded = torch_nn_functional.pad(
|
|
normalized.unsqueeze(1),
|
|
(blur_radius, blur_radius, blur_radius, blur_radius),
|
|
mode="replicate",
|
|
)
|
|
blurred = torch_nn_functional.avg_pool2d(
|
|
padded,
|
|
kernel_size=kernel_size,
|
|
stride=1,
|
|
)
|
|
return blurred.squeeze(1).clamp(0.0, 1.0).to(dtype=torch.float32)
|
|
|
|
|
|
def fill_mask_holes(mask):
|
|
normalized = normalize_mask(mask)
|
|
if np is None:
|
|
return normalized
|
|
|
|
binary = (normalized > 0.5).detach().cpu().numpy().astype(np.uint8)
|
|
filled = np.zeros_like(binary, dtype=np.float32)
|
|
for index, sample in enumerate(binary):
|
|
filled[index] = _fill_holes_binary(sample)
|
|
|
|
return torch.from_numpy(filled).to(device=normalized.device, dtype=torch.float32).clamp(0.0, 1.0)
|
|
|
|
|
|
def remove_small_mask_regions(mask, min_area):
|
|
normalized = normalize_mask(mask)
|
|
area_threshold = _coerce_non_negative_int(min_area)
|
|
if area_threshold <= 1 or np is None:
|
|
return normalized
|
|
|
|
binary = (normalized > 0.5).detach().cpu().numpy().astype(np.uint8)
|
|
cleaned = np.zeros_like(binary, dtype=np.float32)
|
|
for index, sample in enumerate(binary):
|
|
cleaned[index] = _remove_small_regions_binary(sample, area_threshold)
|
|
|
|
return torch.from_numpy(cleaned).to(device=normalized.device, dtype=torch.float32).clamp(0.0, 1.0)
|
|
|
|
|
|
def resolve_image_size(image):
|
|
shape = tuple(getattr(image, "shape", ()))
|
|
if len(shape) != 4:
|
|
return None
|
|
batch = max(1, int(shape[0]))
|
|
height = max(1, int(shape[1]))
|
|
width = max(1, int(shape[2]))
|
|
return batch, height, width
|
|
|
|
|
|
def _fill_holes_binary(sample):
|
|
height, width = sample.shape
|
|
visited = np.zeros((height, width), dtype=bool)
|
|
queue = deque()
|
|
|
|
def enqueue(y, x):
|
|
if 0 <= y < height and 0 <= x < width and sample[y, x] == 0 and not visited[y, x]:
|
|
visited[y, x] = True
|
|
queue.append((y, x))
|
|
|
|
for x in range(width):
|
|
enqueue(0, x)
|
|
enqueue(height - 1, x)
|
|
for y in range(height):
|
|
enqueue(y, 0)
|
|
enqueue(y, width - 1)
|
|
|
|
while queue:
|
|
y, x = queue.popleft()
|
|
enqueue(y - 1, x)
|
|
enqueue(y + 1, x)
|
|
enqueue(y, x - 1)
|
|
enqueue(y, x + 1)
|
|
|
|
result = sample.copy()
|
|
result[(sample == 0) & (~visited)] = 1
|
|
return result.astype(np.float32)
|
|
|
|
|
|
def _remove_small_regions_binary(sample, min_area):
|
|
height, width = sample.shape
|
|
visited = np.zeros((height, width), dtype=bool)
|
|
result = np.zeros((height, width), dtype=np.float32)
|
|
|
|
for start_y in range(height):
|
|
for start_x in range(width):
|
|
if sample[start_y, start_x] == 0 or visited[start_y, start_x]:
|
|
continue
|
|
|
|
queue = deque([(start_y, start_x)])
|
|
visited[start_y, start_x] = True
|
|
region = []
|
|
|
|
while queue:
|
|
y, x = queue.popleft()
|
|
region.append((y, x))
|
|
for next_y, next_x in (
|
|
(y - 1, x),
|
|
(y + 1, x),
|
|
(y, x - 1),
|
|
(y, x + 1),
|
|
):
|
|
if (
|
|
0 <= next_y < height
|
|
and 0 <= next_x < width
|
|
and sample[next_y, next_x] == 1
|
|
and not visited[next_y, next_x]
|
|
):
|
|
visited[next_y, next_x] = True
|
|
queue.append((next_y, next_x))
|
|
|
|
if len(region) >= min_area:
|
|
for y, x in region:
|
|
result[y, x] = 1.0
|
|
|
|
return result
|