255 lines
9.9 KiB
Python
255 lines
9.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""图像预处理 / 后处理。
|
||
|
||
所有数值行为对齐 ComfyUI 节点 ``comfyui_birefnet_ll``：
|
||
* 预处理 —— 与 birefnetNode.ImagePreprocessor 相同的 Resize + ImageNet Normalize
|
||
* 遮罩还原 —— 等价 comfy.utils.common_upscale（F.interpolate，无抗锯齿）
|
||
* 前景精修 —— fast-foreground-estimation（Photoroom）的 box-blur 版本
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from functools import lru_cache
|
||
from typing import Dict, Iterable, Optional, Sequence, Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image, ImageOps
|
||
|
||
# 与推理无关的通用工具（后缀集合、文件名清洗、尺寸换算）放在 fsutil 里，
|
||
# 这里重新导出，保持 `imageops.safe_stem` 这类既有调用点不变。
|
||
from .fsutil import ( # noqa: F401
|
||
ALPHA_SUFFIXES,
|
||
IMAGE_SUFFIXES,
|
||
fit_longest_side,
|
||
safe_stem,
|
||
)
|
||
|
||
try: # cv2 的 box blur 速度更快、背景更纯（与节点一致）
|
||
import cv2
|
||
|
||
_HAS_CV2 = True
|
||
except Exception: # pragma: no cover
|
||
cv2 = None # type: ignore[assignment]
|
||
_HAS_CV2 = False
|
||
|
||
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
||
IMAGENET_STD = [0.229, 0.224, 0.225]
|
||
|
||
#: 与节点 interpolation_modes_mapping 对齐（torchvision InterpolationMode 的数值）
|
||
INTERPOLATIONS: Dict[str, int] = {
|
||
"nearest": 0,
|
||
"bilinear": 2,
|
||
"bicubic": 3,
|
||
"nearest-exact": 0,
|
||
}
|
||
|
||
#: 输出遮罩还原时允许的 F.interpolate 模式
|
||
UPSCALE_METHODS = ("bilinear", "nearest", "nearest-exact", "bicubic")
|
||
|
||
#: 预处理尺寸必须是 32 的倍数（Swin 骨干下采样 32 倍）
|
||
SIZE_ALIGN = 32
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 基础转换
|
||
# --------------------------------------------------------------------------- #
|
||
def pil_to_rgb_array(image: Image.Image) -> np.ndarray:
|
||
"""PIL -> uint8 RGB 数组，同时按 EXIF 方向自动旋转。"""
|
||
if image.mode == "RGBA":
|
||
# 有透明通道时先合成到白底，避免 alpha 变黑
|
||
bg = Image.new("RGBA", image.size, (255, 255, 255, 255))
|
||
image = Image.alpha_composite(bg, image)
|
||
elif image.mode not in ("RGB", "L"):
|
||
image = image.convert("RGB")
|
||
image = ImageOps.exif_transpose(image)
|
||
if image.mode != "RGB":
|
||
image = image.convert("RGB")
|
||
return np.asarray(image, dtype=np.uint8)
|
||
|
||
|
||
def to_pil_rgba(rgb: np.ndarray, alpha: Optional[np.ndarray] = None) -> Image.Image:
|
||
"""uint8 RGB (+ float alpha[0,1]) -> PIL 图片。"""
|
||
if alpha is None:
|
||
return Image.fromarray(np.ascontiguousarray(rgb), "RGB")
|
||
a8 = np.clip(np.rint(alpha * 255.0), 0, 255).astype(np.uint8)
|
||
rgba = np.dstack([rgb, a8])
|
||
return Image.fromarray(np.ascontiguousarray(rgba), "RGBA")
|
||
|
||
|
||
def to_pil_mask(alpha: np.ndarray) -> Image.Image:
|
||
"""float mask[0,1] -> 8bit 灰度图。"""
|
||
m8 = np.clip(np.rint(alpha * 255.0), 0, 255).astype(np.uint8)
|
||
return Image.fromarray(np.ascontiguousarray(m8), "L")
|
||
|
||
|
||
def parse_color(value: object, default: Tuple[int, int, int] = (255, 255, 255)) -> Tuple[int, int, int]:
|
||
"""解析 ``#rgb`` / ``#rrggbb`` / ``(r,g,b)`` / int 为 RGB 三元组。"""
|
||
if value is None:
|
||
return default
|
||
if isinstance(value, (list, tuple)) and len(value) >= 3:
|
||
return tuple(int(np.clip(v, 0, 255)) for v in value[:3]) # type: ignore[return-value]
|
||
if isinstance(value, int):
|
||
return ((value >> 16) & 0xFF, (value >> 8) & 0xFF, value & 0xFF)
|
||
if isinstance(value, str):
|
||
s = value.strip().lstrip("#")
|
||
if len(s) == 3:
|
||
s = "".join(ch * 2 for ch in s)
|
||
if len(s) == 6:
|
||
try:
|
||
v = int(s, 16)
|
||
return ((v >> 16) & 0xFF, (v >> 8) & 0xFF, v & 0xFF)
|
||
except ValueError:
|
||
pass
|
||
return default
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 输入尺寸
|
||
# --------------------------------------------------------------------------- #
|
||
def _align(v: float, align: int = SIZE_ALIGN, minimum: int = SIZE_ALIGN) -> int:
|
||
v = int(round(v / align) * align)
|
||
return max(minimum, v)
|
||
|
||
|
||
def build_input_size(
|
||
src_h: int,
|
||
src_w: int,
|
||
mode: str = "square",
|
||
width: int = 1024,
|
||
height: int = 1024,
|
||
longest_side: int = 1024,
|
||
) -> Tuple[int, int]:
|
||
"""计算网络输入尺寸 (h, w)。
|
||
|
||
mode:
|
||
square 固定 1024x1024（与节点默认一致，也是官方推荐分辨率）
|
||
longest 保持宽高比，长边 = longest_side，边长对齐到 32 的倍数
|
||
custom 使用自定义 width / height
|
||
"""
|
||
if src_h <= 0 or src_w <= 0:
|
||
return max(SIZE_ALIGN, int(height)), max(SIZE_ALIGN, int(width))
|
||
|
||
if mode == "custom":
|
||
h, w = _align(int(height)), _align(int(width))
|
||
elif mode == "longest":
|
||
target = max(SIZE_ALIGN, int(longest_side))
|
||
scale = target / float(max(src_h, src_w))
|
||
h, w = _align(src_h * scale), _align(src_w * scale)
|
||
else: # square
|
||
h, w = _align(int(height)), _align(int(width))
|
||
return h, w
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 预处理
|
||
# --------------------------------------------------------------------------- #
|
||
@lru_cache(maxsize=32)
|
||
def _transform(size_hw: Tuple[int, int], method: str):
|
||
"""缓存 torchvision 变换（与节点 ImagePreprocessor 完全一致）。"""
|
||
from torchvision import transforms
|
||
|
||
interp = INTERPOLATIONS.get(method, 2)
|
||
return transforms.Compose(
|
||
[
|
||
transforms.Resize(size_hw, interpolation=interp),
|
||
transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
|
||
]
|
||
)
|
||
|
||
|
||
def preprocess(rgb: np.ndarray, size_hw: Tuple[int, int], method: str = "bilinear") -> torch.Tensor:
|
||
"""uint8 RGB (h,w,3) -> 归一化 float32 张量 (1,3,H,W)。"""
|
||
arr = np.asarray(rgb, dtype=np.float32) / 255.0 # 新建可写数组，避免 from_numpy 告警
|
||
tensor = torch.from_numpy(np.ascontiguousarray(arr)).permute(2, 0, 1).unsqueeze(0)
|
||
return _transform(tuple(size_hw), method)(tensor)
|
||
|
||
|
||
def upscale_mask(mask_bchw: torch.Tensor, height: int, width: int, method: str = "bilinear") -> torch.Tensor:
|
||
"""把网络输出的低分辨率遮罩还原到原图尺寸（等价 comfy.utils.common_upscale）。"""
|
||
mode = method if method in UPSCALE_METHODS else "bilinear"
|
||
if tuple(mask_bchw.shape[-2:]) == (height, width):
|
||
return mask_bchw
|
||
return torch.nn.functional.interpolate(mask_bchw, size=(height, width), mode=mode)
|
||
|
||
|
||
def filter_mask(mask: torch.Tensor, threshold: float) -> torch.Tensor:
|
||
"""低于阈值的概率直接置零（与节点 util.filter_mask 一致）。"""
|
||
if threshold <= 0:
|
||
return mask
|
||
return mask * (mask > threshold).to(mask.dtype)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 前景精修（fast-foreground-estimation）
|
||
# --------------------------------------------------------------------------- #
|
||
def _box_blur(arr: np.ndarray, r: int) -> np.ndarray:
|
||
"""box blur；r 为核尺寸。cv2 可用时优先使用，保证与节点结果一致。
|
||
|
||
注意：cv2 对 (h, w, 1) 的输入会返回 (h, w)，这里统一补齐通道维，
|
||
否则后续与 (h, w, 3) 广播会报错（原节点也是靠 `[:, :, None]` 兜住这一点）。
|
||
"""
|
||
r = int(max(1, r))
|
||
limit = 2 * min(arr.shape[0], arr.shape[1]) + 1
|
||
r = int(min(r, limit))
|
||
if _HAS_CV2:
|
||
out = cv2.blur(np.ascontiguousarray(arr, dtype=np.float32), (r, r))
|
||
return out[..., None] if (out.ndim == 2 and arr.ndim == 3) else out
|
||
# 兜底：torchvision 高斯模糊
|
||
from torchvision.transforms import functional as TF
|
||
|
||
if r % 2 == 0:
|
||
r += 1
|
||
t = torch.from_numpy(np.ascontiguousarray(arr)).permute(2, 0, 1).unsqueeze(0).float()
|
||
out = TF.gaussian_blur(t, r) if r > 1 else t
|
||
return out.squeeze(0).permute(1, 2, 0).numpy()
|
||
|
||
|
||
def _fb_step(image: np.ndarray, F: np.ndarray, B: np.ndarray, alpha: np.ndarray, r: int):
|
||
a = alpha[:, :, None] if alpha.ndim == 2 else alpha
|
||
blurred_alpha = _box_blur(a, r)
|
||
blurred_F = _box_blur(F * a, r) / (blurred_alpha + 1e-5)
|
||
blurred_B = _box_blur(B * (1.0 - a), r) / ((1.0 - blurred_alpha) + 1e-5)
|
||
out = blurred_F + a * (image - a * blurred_F - (1.0 - a) * blurred_B)
|
||
return np.clip(out, 0.0, 1.0), blurred_B
|
||
|
||
|
||
def refine_foreground(rgb01: np.ndarray, alpha: np.ndarray, r1: int = 90, r2: int = 6) -> np.ndarray:
|
||
"""估计前景色，消除半透明边缘残留的原背景色。
|
||
|
||
参考 https://github.com/Photoroom/fast-foreground-estimation
|
||
"""
|
||
image = np.ascontiguousarray(rgb01, dtype=np.float32)
|
||
a = alpha.astype(np.float32)
|
||
F, blur_B = _fb_step(image, image, image, a, r1)
|
||
F2, _ = _fb_step(image, F, blur_B, a, r2)
|
||
return np.clip(F2, 0.0, 1.0)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 合成输出
|
||
# --------------------------------------------------------------------------- #
|
||
def compose_on_color(rgb: np.ndarray, alpha: np.ndarray, color: Sequence[int]) -> Image.Image:
|
||
"""按遮罩把前景合成到纯色背景上，返回 RGB 图。"""
|
||
a = alpha[:, :, None].astype(np.float32)
|
||
bg = np.array(color, dtype=np.float32).reshape(1, 1, 3)
|
||
out = rgb.astype(np.float32) * a + bg * (1.0 - a)
|
||
return Image.fromarray(np.clip(np.rint(out), 0, 255).astype(np.uint8), "RGB")
|
||
|
||
|
||
def save_image(image: Image.Image, path: str, quality: int = 95) -> str:
|
||
"""保存图片（自动建目录）。PNG/WebP 支持透明，JPEG 不支持。"""
|
||
os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
|
||
ext = os.path.splitext(path)[1].lower()
|
||
if ext in (".jpg", ".jpeg"):
|
||
if image.mode == "RGBA":
|
||
bg = Image.new("RGBA", image.size, (255, 255, 255, 255))
|
||
image = Image.alpha_composite(bg, image).convert("RGB")
|
||
image.save(path, quality=quality, subsampling=0, optimize=True)
|
||
elif ext == ".webp":
|
||
image.save(path, quality=quality, method=4)
|
||
else:
|
||
image.save(path, optimize=False)
|
||
return path
|