rui40000andClaude Opus 4.8
1bd90c71b6
满屏文字水印（watermark_node.py）: - 文案/字体/字号/旋转角度/密度/字间距/透明度/颜色 八项可调 - 交错网格平铺后整体旋转再从中心裁切，任意角度下四角均无空白 FeyNobg 抠图（feynobg_node.py + feynobg/）: - 内嵌 nobg 推理子集（Apache-2.0），全自动去背景，输出 alpha 与去背景图 - 重写预处理去掉对 transformers>=5.4 的依赖，4.x 环境可直接使用 - 修复权重键名与 transformers 4.x 的 SwinBackbone 命名不兼容： 不处理时 958 个参数仅 405 个对得上，backbone 形同随机初始化， 模型不报错但 alpha 几乎全黑；现按环境自动重映射并严格校验， 除确定性 buffer 外任何失配都直接中止 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def iou_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
inter = (pred * target).sum(dim=(2, 3))
|
|
union = pred.sum(dim=(2, 3)) + target.sum(dim=(2, 3)) - inter
|
|
return (1 - inter / (union + 1e-8)).mean()
|
|
|
|
|
|
def ssim_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
C1 = 0.01**2
|
|
C2 = 0.03**2
|
|
mu_x = F.avg_pool2d(pred, 3, 1, 1)
|
|
mu_y = F.avg_pool2d(target, 3, 1, 1)
|
|
sigma_x = F.avg_pool2d(pred * pred, 3, 1, 1) - mu_x * mu_x
|
|
sigma_y = F.avg_pool2d(target * target, 3, 1, 1) - mu_y * mu_y
|
|
sigma_xy = F.avg_pool2d(pred * target, 3, 1, 1) - mu_x * mu_y
|
|
ssim_map = ((2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2)) / (
|
|
(mu_x * mu_x + mu_y * mu_y + C1) * (sigma_x + sigma_y + C2)
|
|
)
|
|
return torch.clamp((1 - ssim_map) / 2, 0, 1).mean()
|
|
|
|
|
|
def birefnet_loss(scaled_preds: list[torch.Tensor], gt: torch.Tensor) -> torch.Tensor:
|
|
"""Multi-scale pixel loss matching BiRefNet training: weighted BCE + IoU + SSIM."""
|
|
loss = torch.tensor(0.0, device=gt.device)
|
|
for pred in scaled_preds:
|
|
if pred.shape[2:] != gt.shape[2:]:
|
|
pred = F.interpolate(
|
|
pred, size=gt.shape[2:], mode="bilinear", align_corners=True
|
|
)
|
|
pred_sig = pred.sigmoid()
|
|
loss = loss + 30 * F.binary_cross_entropy_with_logits(pred, gt)
|
|
loss = loss + 0.5 * iou_loss(pred_sig, gt)
|
|
loss = loss + 10 * ssim_loss(pred_sig, gt)
|
|
return loss
|