306 lines
13 KiB
Python
306 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
merge_fixed.py — 原图 + AI修改图 无缝融合工具
|
||
============================================
|
||
背景：ComfyUI inpainting 移除合影中的其他人物时，即使指定"非修改区域不得像素漂移"，
|
||
全图清晰度仍不可避免损失。本工具把：
|
||
- 原图的"未修改区域"（清晰度完好）
|
||
- 修改图的"修改区域"（人物被干净移除）
|
||
自动合并，过渡自然无痕迹。
|
||
|
||
用法：
|
||
python merge_fixed.py <原图目录> <修改图目录> <输出目录> [--feather 25] [--threshold 20]
|
||
|
||
匹配规则：
|
||
- 原图 IMG_0463.jpg（任意后缀 jpg/png/webp/jpeg）
|
||
- 修改图 IMG_0463*.png（以原图文件名开头、.png 结尾、中间任意字符）
|
||
- 一张原图可有多张修改图，每张都独立合并
|
||
- 输出：{原图stem}_{修改图中间部分}.png
|
||
|
||
算法：
|
||
1. 修改图 resize 到原图尺寸
|
||
2. Lab 色彩空间差异图 → 高斯模糊去噪 → 阈值 → 闭运算+膨胀 = 修改区 mask
|
||
3. 羽化混合：result = 原图*(1-alpha) + 修改图*alpha（alpha 由 mask 高斯模糊得到）
|
||
"""
|
||
import argparse
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||
|
||
# 复用 face_checker 的人脸检测（目标人物面部保护）
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||
import face_checker as fc
|
||
|
||
IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
|
||
_fc_det = None
|
||
|
||
|
||
def get_face_detector():
|
||
global _fc_det
|
||
if _fc_det is None:
|
||
_fc_det = fc.FaceDetector()
|
||
return _fc_det
|
||
|
||
|
||
def protect_faces(alpha, mod_rgb, expand=1.35, margin=40):
|
||
"""
|
||
面部强制保护（保险）：检测修改图中残留的人脸（=目标人物，被移除者已被 inpaint 掉），
|
||
把这些人脸区域的 alpha 置 0（100% 用原图像素，零改变）。
|
||
通用规则（有限修改→原图）已覆盖大部分，此项兜底确保面部绝对不变。
|
||
"""
|
||
det = get_face_detector()
|
||
faces = fc._detect_faces_fast(mod_rgb, det)
|
||
if not faces:
|
||
return alpha
|
||
for f in faces:
|
||
x, y, w, h = f[0], f[1], f[2], f[3]
|
||
cw, ch = w * expand, h * expand
|
||
x0 = max(0, int(x - (cw - w) / 2) - margin)
|
||
y0 = max(0, int(y - (ch - h) / 2) - margin)
|
||
x1 = min(alpha.shape[1], int(x + w + (cw - w) / 2) + margin)
|
||
y1 = min(alpha.shape[0], int(y + h + (ch - h) / 2) + margin)
|
||
alpha[y0:y1, x0:x1] = 0.0
|
||
return alpha
|
||
|
||
|
||
def build_alpha_from_mask(mask_rgb, feather=20):
|
||
"""
|
||
从 ComfyUI inpaint 蒙版生成 alpha：白色区域 = 被修改区（100% 修改图），黑色 = 原图。
|
||
蒙版是最准确的"修改区"标记（用户在 ComfyUI 画的就是要移除的人物），
|
||
比 diff 检测可靠 100 倍——紧贴人物的残影完美解决。
|
||
feather: 蒙版边缘羽化像素
|
||
"""
|
||
gray = cv2.cvtColor(mask_rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
|
||
# 蒙版白色→1（修改图），黑色→0（原图）
|
||
alpha = gray
|
||
if feather > 0:
|
||
k = feather * 2 + 1
|
||
alpha = cv2.GaussianBlur(alpha, (k, k), 0)
|
||
return alpha
|
||
|
||
|
||
def find_mask(orig_path, mask_dir):
|
||
"""按文件名匹配蒙版：{stem}*mask*.png 或 {stem}*.png（取 mask 关键字优先）"""
|
||
stem = Path(orig_path).stem
|
||
cands = sorted(p for p in mask_dir.iterdir()
|
||
if p.suffix.lower() == ".png" and p.name.startswith(stem))
|
||
# 优先带 mask 关键字的
|
||
for p in cands:
|
||
if "mask" in p.name.lower():
|
||
return p
|
||
return cands[0] if cands else None
|
||
|
||
|
||
def build_alpha(orig_rgb, mod_rgb, blur_kernel=15, threshold=4.0, min_area_ratio=0.005,
|
||
softness=1.5, dilate=60, body_extend=2.2):
|
||
"""
|
||
融合权重 alpha（0=原图像素, 1=修改图像素）——双图人脸差集定位被移除人物：
|
||
1. 原图人脸 - 修改图人脸 = 被移除人物（合影中被清掉的人）
|
||
2. 被移除人物区域（脸框 + 向下 body_extend 倍覆盖身体）= 修改区
|
||
3. 目标人物（修改图残留的最大人脸）紧贴框 = 保留区（强制原图）
|
||
4. diff 高值连通块（排除保留区）兜底其他修改
|
||
5. 膨胀覆盖边缘残影
|
||
紧贴合影场景：被移除人物与目标紧贴时，靠"人脸差集"精确定位，不靠 diff 幅度。
|
||
"""
|
||
det = get_face_detector()
|
||
orig_faces = fc._detect_faces_fast(orig_rgb, det)
|
||
mod_faces = fc._detect_faces_fast(mod_rgb, det)
|
||
h, w = orig_rgb.shape[:2]
|
||
|
||
def matched(of, mfaces):
|
||
ox, oy = of[0] + of[2] / 2, of[1] + of[3] / 2
|
||
return any(abs(ox - (mf[0] + mf[2] / 2)) < of[2] * 0.8 and abs(oy - (mf[1] + mf[3] / 2)) < of[3] * 0.8
|
||
for mf in mfaces)
|
||
|
||
removed_faces = [f for f in orig_faces if not matched(f, mod_faces)]
|
||
# 过滤误检小脸（背景物体/花纹被当脸）：只保留足够大的人脸
|
||
# 真实人物脸 ≥ 图面积 0.05% 或短边 ≥ 40px
|
||
min_face_area = max(1, int(h * w * 0.0005))
|
||
removed_faces = [f for f in removed_faces
|
||
if f[2] * f[3] >= min_face_area and min(f[2], f[3]) >= 40]
|
||
|
||
# 被移除人物区（脸 + 身体向下扩展）
|
||
mod_mask = np.zeros((h, w), np.uint8)
|
||
for f in removed_faces:
|
||
x, y, fw, fh = [int(v) for v in f[:4]]
|
||
x0 = max(0, int(x - fw * 0.4)); x1 = min(w, int(x + fw * 1.4))
|
||
y0 = max(0, int(y - fh * 0.3)); y1 = min(h, int(y + fh * body_extend))
|
||
mod_mask[y0:y1, x0:x1] = 1
|
||
|
||
# 保留区：目标人物最大脸，紧贴框（左右不扩，防侵入紧贴的被移除人物）
|
||
keep_mask = np.zeros((h, w), np.uint8)
|
||
if mod_faces:
|
||
big = max(mod_faces, key=lambda f: f[2] * f[3])
|
||
x, y, fw, fh = [int(v) for v in big[:4]]
|
||
x0, x1 = max(0, x), min(w, x + fw)
|
||
y0, y1 = max(0, int(y - fh * 0.35)), min(h, int(y + fh * 1.2))
|
||
keep_mask[y0:y1, x0:x1] = 1
|
||
|
||
# diff 兜底连通块（排除保留区）。无被移除人物时用高阈值（保守，防降质误判）
|
||
o_lab = cv2.cvtColor(orig_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
|
||
m_lab = cv2.cvtColor(mod_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
|
||
diff_s = cv2.GaussianBlur(np.abs(o_lab - m_lab).mean(axis=2), (blur_kernel, blur_kernel), 0)
|
||
eff_threshold = threshold if removed_faces else max(threshold, 15.0)
|
||
base = ((diff_s > eff_threshold) & ~keep_mask.astype(bool)).astype(np.uint8)
|
||
num, labels, stats, _ = cv2.connectedComponentsWithStats(base, 8)
|
||
min_area = max(1, int(min_area_ratio * h * w))
|
||
filtered = np.zeros_like(base)
|
||
for i in range(1, num):
|
||
if stats[i, cv2.CC_STAT_AREA] >= min_area:
|
||
filtered[labels == i] = 1
|
||
|
||
core = np.maximum(filtered, mod_mask)
|
||
if dilate > 0:
|
||
core = cv2.dilate(core, np.ones((dilate, dilate), np.uint8))
|
||
alpha = core.astype(np.float32)
|
||
if softness > 0:
|
||
k = int(softness * 6) * 2 + 1
|
||
alpha = cv2.GaussianBlur(alpha, (k, k), 0)
|
||
return alpha
|
||
"""
|
||
融合权重 alpha（0=原图像素, 1=修改图像素）：
|
||
- 固定阈值判定明显修改区（人物移除 diff 高，有限降质 diff 低，区分度 ~20x）
|
||
- diff > 阈值 → alpha≈1（100% 修改图，修改区内部零残影）
|
||
- diff < 阈值 → alpha≈0（100% 原图，目标人物/有限修改区零改变）
|
||
- 修改区膨胀 dilate 像素（默认 75）：覆盖人物边缘 10-40px 的浅残影带
|
||
（残影带 diff 接近未修改区，无法用阈值检测，必须靠膨胀）
|
||
- 陡 sigmoid 边缘过渡防硬边
|
||
"""
|
||
o = cv2.cvtColor(orig_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
|
||
m = cv2.cvtColor(mod_rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
|
||
diff = np.abs(o - m).mean(axis=2)
|
||
diff_s = cv2.GaussianBlur(diff, (blur_kernel, blur_kernel), 0)
|
||
|
||
alpha = 1.0 / (1.0 + np.exp(-(diff_s - threshold) / softness))
|
||
|
||
# 修改区膨胀：覆盖人物边缘残影带
|
||
if dilate > 0:
|
||
kernel = np.ones((dilate, dilate), np.uint8)
|
||
core = (alpha > 0.5).astype(np.uint8)
|
||
core = cv2.dilate(core, kernel)
|
||
alpha = np.maximum(alpha, core.astype(np.float32))
|
||
return alpha
|
||
|
||
|
||
def seamless_merge(orig_rgb, mod_rgb, alpha):
|
||
"""
|
||
连续 alpha 混合：result = orig*(1-alpha) + mod*alpha。
|
||
alpha 已是平滑的浮点图（sigmoid 过渡），无需额外羽化。
|
||
"""
|
||
a = alpha[..., None].astype(np.float32)
|
||
result = orig_rgb.astype(np.float32) * (1.0 - a) + mod_rgb.astype(np.float32) * a
|
||
return np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
|
||
def find_modified(orig_path, mod_dir):
|
||
"""按文件名匹配修改图：{stem}*.png"""
|
||
stem = Path(orig_path).stem
|
||
return sorted(p for p in mod_dir.iterdir()
|
||
if p.suffix.lower() == ".png" and p.name.startswith(stem))
|
||
|
||
|
||
def process(orig_dir, mod_dir, out_dir, alpha_threshold=8.0, softness=1.5, protect=True,
|
||
mask_dir=None, feather=20, verbose=True):
|
||
"""
|
||
alpha_threshold: 无蒙版时的 diff 阈值（默认 8；误判时调）
|
||
softness: 过渡带宽度
|
||
protect: 面部强制保护（无蒙版 fallback 时的保险）
|
||
mask_dir: ComfyUI inpaint 蒙版目录（可选，有则优先用，蒙版白色=被移除区，最准确）
|
||
feather: 蒙版边缘羽化
|
||
"""
|
||
orig_dir, mod_dir, out_dir = Path(orig_dir), Path(mod_dir), Path(out_dir)
|
||
if mask_dir:
|
||
mask_dir = Path(mask_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
originals = sorted(p for p in orig_dir.iterdir()
|
||
if p.suffix.lower() in IMG_EXTS and not p.name.startswith("."))
|
||
|
||
if not originals:
|
||
print(f"[WARN] 原图目录没有图片: {orig_dir}")
|
||
return
|
||
|
||
total = 0
|
||
for op in originals:
|
||
mods = find_modified(op, mod_dir)
|
||
if not mods:
|
||
if verbose:
|
||
print(f"[跳过] {op.name}: 无匹配修改图")
|
||
continue
|
||
|
||
from PIL import Image
|
||
orig_img = Image.open(op).convert("RGB")
|
||
ow, oh = orig_img.size
|
||
orig_rgb = np.array(orig_img)
|
||
|
||
# 尝试匹配蒙版
|
||
mask_path = find_mask(op, mask_dir) if mask_dir else None
|
||
|
||
for mp in mods:
|
||
try:
|
||
mod_img = Image.open(mp).convert("RGB")
|
||
if mod_img.size != (ow, oh):
|
||
mod_img = mod_img.resize((ow, oh), Image.LANCZOS)
|
||
mod_rgb = np.array(mod_img)
|
||
|
||
if mask_path and mask_path.exists():
|
||
# 蒙版模式：最准确
|
||
mask_img = Image.open(mask_path).convert("RGB")
|
||
if mask_img.size != (ow, oh):
|
||
mask_img = mask_img.resize((ow, oh), Image.LANCZOS)
|
||
mask_rgb = np.array(mask_img)
|
||
alpha = build_alpha_from_mask(mask_rgb, feather=feather)
|
||
mode = f"蒙版模式({mask_path.name})"
|
||
else:
|
||
# 双图人脸差集模式（自动定位被移除人物 + 保留目标脸）
|
||
alpha = build_alpha(orig_rgb, mod_rgb, threshold=alpha_threshold, softness=softness)
|
||
mode = "差集模式"
|
||
|
||
area_ratio = float((alpha > 0.5).mean())
|
||
if area_ratio < 0.0005:
|
||
if verbose:
|
||
print(f"[警告] {op.name} <- {mp.name}: 明显修改区仅 {area_ratio*100:.2f}%%，可能未检测到修改")
|
||
mask_area_note = f"明显修改区 {area_ratio*100:.2f}% (偏小?)"
|
||
else:
|
||
mask_area_note = f"明显修改区 {area_ratio*100:.1f}%"
|
||
|
||
result = seamless_merge(orig_rgb, mod_rgb, alpha)
|
||
mid = mp.stem[len(Path(op).stem):] or "_mod"
|
||
out_name = f"{Path(op).stem}{mid}.png"
|
||
Image.fromarray(result).save(out_dir / out_name)
|
||
total += 1
|
||
if verbose:
|
||
print(f"[OK] {out_name} | {mode} | {mask_area_note}")
|
||
except Exception as e:
|
||
if verbose:
|
||
print(f"[失败] {op.name} <- {mp.name}: {e}")
|
||
|
||
print(f"\n完成：共输出 {total} 张融合图 -> {out_dir}")
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="原图 + AI修改图 无缝融合工具（支持 ComfyUI 蒙版）")
|
||
ap.add_argument("orig_dir", help="原图目录")
|
||
ap.add_argument("mod_dir", help="修改图目录（文件名以原图名开头、.png 结尾）")
|
||
ap.add_argument("out_dir", help="输出目录")
|
||
ap.add_argument("--mask-dir", default=None,
|
||
help="ComfyUI inpaint 蒙版目录（可选，有则优先用；蒙版白色=被移除区，最准确解决紧贴残影）")
|
||
ap.add_argument("--alpha-threshold", type=float, default=8.0,
|
||
help="无蒙版时的 diff 阈值（默认 8）")
|
||
ap.add_argument("--softness", type=float, default=1.5, help="过渡带宽度（默认 1.5）")
|
||
ap.add_argument("--feather", type=int, default=20, help="蒙版边缘羽化（默认 20）")
|
||
ap.add_argument("--no-protect-face", action="store_true", help="关闭面部强制保护（无蒙版时的保险）")
|
||
args = ap.parse_args()
|
||
|
||
process(args.orig_dir, args.mod_dir, args.out_dir,
|
||
alpha_threshold=args.alpha_threshold, softness=args.softness,
|
||
protect=not args.no_protect_face, mask_dir=args.mask_dir, feather=args.feather)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|