rui40000andClaude Opus 4.8
be93fd36a3
针对「主体整片半透明发灰」的实际反馈，模型侧无可调参数 （架构参数训练时已固化），故在后处理层开放： - alpha_threshold / alpha_softness：以阈值为中心、柔和度为宽度 取一段区间线性拉伸到 [0,1]，压掉灰雾同时保留边缘过渡。 默认 (0.5, 1.0) 区间恰为 [0,1]，是恒等变换，行为与之前逐位一致； 实测 (0.35, 0.3) 半透明像素占比 1.49% -> 0.31%，主体均值几乎不变 - keep_aspect_ratio：模型固定吃 1024 方图，长图默认被拉伸压扁。 开启后改为等比缩放 + 边缘延展补边（不填黑，避免造出假轮廓）， 推理后按有效区域裁掉补边 同步更新示例工作流的 widgets_values（widget 由 5 个增至 8 个， 不同步会导致参数错位）与 Note 说明。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
167 lines
7.4 KiB
Python
167 lines
7.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
BiRefNet 预处理 / 后处理（Ruinode 内嵌版）
|
||
==========================================
|
||
官方 nobg 的同名文件继承 transformers>=5.4 的 TorchvisionBackend，而 ComfyUI
|
||
常见环境仍是 transformers 4.x（本机 T8 为 4.49.0），直接引入会报：
|
||
|
||
ModuleNotFoundError: No module named 'transformers.image_processing_backends'
|
||
|
||
升级 transformers 到 5.x 极可能连累其他自定义节点，因此这里用纯 torch/PIL
|
||
复刻官方的同一套数值规格，去掉对 transformers 的依赖：
|
||
|
||
- 预处理：转 RGB → 双线性缩放到 size（抗锯齿）→ 归一到 [0,1] → ImageNet 标准化
|
||
- 后处理：logits.sigmoid() → 双线性缩放回原尺寸
|
||
（**先 sigmoid 再缩放**，与官方 post_process_alpha_matting 及其 eval 脚本一致，
|
||
顺序颠倒会让边缘过渡发生偏移）
|
||
- cutout：把 alpha 写入原图的 alpha 通道，输出 RGBA
|
||
|
||
除去掉 transformers 依赖外，逐项对齐官方实现；训练用的 segmentation_maps /
|
||
labels 分支推理不需要，未移植。
|
||
"""
|
||
import json
|
||
import os
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
# 与 transformers.image_utils 的 IMAGENET_DEFAULT_MEAN / STD 相同，
|
||
# 也与 FeyNobg 的 preprocessor_config.json 一致
|
||
IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
|
||
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
|
||
|
||
|
||
class BiRefNetImageProcessor:
|
||
"""BiRefNet 图像处理器（预处理 + 后处理）。"""
|
||
|
||
def __init__(self, size=None, image_mean=None, image_std=None,
|
||
rescale_factor=1 / 255, **kwargs):
|
||
size = size or {"height": 1024, "width": 1024}
|
||
self.size = {
|
||
"height": int(size.get("height", 1024)),
|
||
"width": int(size.get("width", 1024)),
|
||
}
|
||
self.image_mean = tuple(image_mean or IMAGENET_DEFAULT_MEAN)
|
||
self.image_std = tuple(image_std or IMAGENET_DEFAULT_STD)
|
||
self.rescale_factor = rescale_factor
|
||
|
||
# ------------------------------------------------------------------ 构造
|
||
@classmethod
|
||
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
|
||
"""从本地目录的 preprocessor_config.json 读取配置（不联网）。"""
|
||
cfg = {}
|
||
if os.path.isdir(pretrained_model_name_or_path):
|
||
p = os.path.join(pretrained_model_name_or_path,
|
||
"preprocessor_config.json")
|
||
if os.path.isfile(p):
|
||
with open(p, "r", encoding="utf-8") as f:
|
||
cfg = json.load(f)
|
||
cfg.update(kwargs)
|
||
return cls(
|
||
size=cfg.get("size"),
|
||
image_mean=cfg.get("image_mean"),
|
||
image_std=cfg.get("image_std"),
|
||
rescale_factor=cfg.get("rescale_factor", 1 / 255),
|
||
)
|
||
|
||
# ---------------------------------------------------------------- 预处理
|
||
def preprocess_tensor(self, images, device=None, dtype=torch.float32,
|
||
keep_aspect=False):
|
||
"""
|
||
ComfyUI 原生张量入口，避免 PIL 往返。
|
||
|
||
images: (B, H, W, 3) float，值域 [0,1]（ComfyUI 的 IMAGE 约定）
|
||
keep_aspect: False 走官方路径，直接拉伸成正方形（与训练一致）；
|
||
True 则按长边等比缩放后补边，保留原始比例。
|
||
|
||
返回 (pixel_values, meta):
|
||
pixel_values (B, 3, size_h, size_w)，已 ImageNet 标准化
|
||
meta keep_aspect 时为有效区域 (nh, nw)，供后处理裁掉补边；否则 None
|
||
"""
|
||
x = images.permute(0, 3, 1, 2).contiguous().float() # BHWC -> BCHW
|
||
th, tw = self.size["height"], self.size["width"]
|
||
meta = None
|
||
|
||
if keep_aspect:
|
||
_, _, H, W = x.shape
|
||
scale = min(th / H, tw / W)
|
||
nh = max(1, min(th, int(round(H * scale))))
|
||
nw = max(1, min(tw, int(round(W * scale))))
|
||
x = F.interpolate(x, size=(nh, nw), mode="bilinear",
|
||
align_corners=False, antialias=True)
|
||
# 补边放在右/下侧，用边缘像素延展而非填黑：填黑会凭空造出一条
|
||
# 高对比直边，模型容易把它当成物体轮廓
|
||
if nh < th or nw < tw:
|
||
x = F.pad(x, (0, tw - nw, 0, th - nh), mode="replicate")
|
||
meta = (nh, nw)
|
||
else:
|
||
# 官方走 torchvision 的 resize，对张量默认开抗锯齿；缩小到 1024 时
|
||
# 是否抗锯齿对细边缘影响可见，这里保持一致
|
||
x = F.interpolate(x, size=(th, tw), mode="bilinear",
|
||
align_corners=False, antialias=True)
|
||
|
||
mean = torch.tensor(self.image_mean, dtype=torch.float32).view(1, 3, 1, 1)
|
||
std = torch.tensor(self.image_std, dtype=torch.float32).view(1, 3, 1, 1)
|
||
x = (x - mean) / std
|
||
x = x.to(device=device, dtype=dtype) if device is not None \
|
||
else x.to(dtype=dtype)
|
||
return x, meta
|
||
|
||
def __call__(self, images, return_tensors="pt", **kwargs):
|
||
"""PIL 入口，兼容官方 README 的用法：processor(image, return_tensors='pt')。"""
|
||
import numpy as np
|
||
from PIL import Image as PILImage
|
||
|
||
if isinstance(images, PILImage.Image):
|
||
images = [images]
|
||
arrs = []
|
||
for im in images:
|
||
im = im.convert("RGB")
|
||
arrs.append(np.asarray(im, dtype=np.float32) * self.rescale_factor)
|
||
batch = torch.from_numpy(np.stack(arrs)) # (B,H,W,3) in [0,1]
|
||
return {"pixel_values": self.preprocess_tensor(batch)[0]}
|
||
|
||
# ---------------------------------------------------------------- 后处理
|
||
def post_process_alpha_matting(self, outputs, target_sizes=None, crop=None):
|
||
"""
|
||
原始 logits -> 每图 alpha（[0,1]，形状 (H, W)）。
|
||
|
||
outputs: 含 "logits" 的 dict 或带 .logits 的对象，logits 形状 (B,1,H,W)
|
||
target_sizes: [(h, w), ...]，逐图缩放回原尺寸
|
||
crop: 预处理若做过等比补边，这里传 (nh, nw) 把补出来的部分裁掉
|
||
"""
|
||
logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
|
||
if crop is not None:
|
||
nh, nw = crop
|
||
logits = logits[:, :, :nh, :nw]
|
||
if target_sizes is not None and len(logits) != len(target_sizes):
|
||
raise ValueError(
|
||
f"给了 {len(target_sizes)} 个目标尺寸，但批次里有 {len(logits)} 张图"
|
||
)
|
||
# 官方注释明确：先 sigmoid 再缩放，与 eval/benchmark 脚本一致
|
||
probs = logits.sigmoid()
|
||
mattes = []
|
||
for idx in range(len(probs)):
|
||
alpha = probs[idx].unsqueeze(0) # (1,1,H,W)
|
||
if target_sizes is not None:
|
||
alpha = F.interpolate(
|
||
alpha, size=target_sizes[idx],
|
||
mode="bilinear", align_corners=False,
|
||
)
|
||
mattes.append(alpha[0, 0])
|
||
return mattes
|
||
|
||
@staticmethod
|
||
def cutout(image, alpha):
|
||
"""把 alpha 合成到原图，返回 RGBA。alpha 可为 (H,W) 张量或 PIL 图。"""
|
||
from PIL import Image as PILImage
|
||
|
||
if isinstance(alpha, torch.Tensor):
|
||
arr = (alpha.detach().clamp(0, 1) * 255).to(torch.uint8).cpu().numpy()
|
||
alpha = PILImage.fromarray(arr, mode="L")
|
||
if alpha.size != image.size:
|
||
alpha = alpha.resize(image.size, PILImage.Resampling.BILINEAR)
|
||
cutout = image.convert("RGBA")
|
||
cutout.putalpha(alpha)
|
||
return cutout
|