dechecker.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. """dechecker.py —— 兜底用:把 AI 图里画死的"透明棋盘格/背景"抠掉,恢复真透明。
  2. 重要:现在生图管线已直接向接口请求真·透明 PNG(providers + pipeline),
  3. 正常情况下根本不需要本模块。只有当上游模型仍把背景画成棋盘格时,才作为
  4. 最后兜底自动运行。
  5. 关键改进(避免戳洞):只移除"与图像边缘相连"的背景(四周的棋盘格 + 透明边),
  6. 角色内部的白色高光/亮点因为被角色包住、不与边缘相连,会被完整保留——
  7. 不会再像以前那样把白色切成缺口。
  8. """
  9. import os
  10. import sys
  11. import numpy as np
  12. from PIL import Image, ImageFilter
  13. def _border_outside(bg):
  14. """返回与图像四边相连的背景区域(布尔)。优先用 cv2,没有则用 numpy 形态学重建。"""
  15. try:
  16. import cv2
  17. n, labels = cv2.connectedComponents(bg.astype(np.uint8), connectivity=4)
  18. border = set(labels[0, :]).union(labels[-1, :], labels[:, 0], labels[:, -1])
  19. border.discard(0)
  20. return np.isin(labels, list(border))
  21. except Exception:
  22. # numpy 兜底:从四边种子在 bg 内反复膨胀直到稳定(形态学重建)
  23. seed = np.zeros_like(bg)
  24. seed[0, :] = bg[0, :]; seed[-1, :] = bg[-1, :]
  25. seed[:, 0] = bg[:, 0]; seed[:, -1] = bg[:, -1]
  26. for _ in range(4000):
  27. grown = seed | np.roll(seed, 1, 0) | np.roll(seed, -1, 0) \
  28. | np.roll(seed, 1, 1) | np.roll(seed, -1, 1)
  29. grown &= bg
  30. if np.array_equal(grown, seed):
  31. break
  32. seed = grown
  33. return seed
  34. def clean(im):
  35. """返回 (新图RGBA, 抹掉的像素数)。无棋盘/已透明的图基本无操作。"""
  36. img = im.convert("RGBA")
  37. arr = np.asarray(img).copy()
  38. al = arr[:, :, 3]
  39. rgb = arr[:, :, :3].astype(np.int16)
  40. bright = rgb.max(2)
  41. sat = bright - rgb.min(2)
  42. light = (al >= 200) & (bright > 185) & (sat < 28) # 棋盘格的白/浅灰方块
  43. trans = al < 40
  44. if light.mean() < 0.04: # 几乎没有浅色块 → 不是棋盘
  45. return img, 0
  46. bg = light | trans
  47. outside = _border_outside(bg)
  48. # 只在"边缘相连的浅色块"占比可观时才动手(避免误伤纯色背景插画)
  49. if (outside & light).mean() < 0.03:
  50. return img, 0
  51. # 向内吃 2px 去掉灰边
  52. m = outside
  53. for _ in range(2):
  54. m = (m | np.roll(m, 1, 0) | np.roll(m, -1, 0)
  55. | np.roll(m, 1, 1) | np.roll(m, -1, 1))
  56. removed = int(((al > 10) & m).sum())
  57. a = al.copy(); a[m] = 0
  58. arr[:, :, 3] = a
  59. arr[m, 0] = 0; arr[m, 1] = 0; arr[m, 2] = 0
  60. out = Image.fromarray(arr.astype(np.uint8), "RGBA")
  61. out.putalpha(out.getchannel("A").filter(ImageFilter.GaussianBlur(1.0))) # 柔边
  62. return out, removed
  63. def declutter_img(img, feather=False):
  64. out, removed = clean(img)
  65. return out.convert("RGBA"), removed > 0
  66. def declutter_file(src, dst=None):
  67. out, removed = clean(Image.open(src))
  68. if removed > 0:
  69. out.save(dst or src)
  70. return removed > 0
  71. def _walk(paths):
  72. for p in paths:
  73. if os.path.isdir(p):
  74. for root, _d, files in os.walk(p):
  75. for f in files:
  76. if f.lower().endswith(".png"):
  77. yield os.path.join(root, f)
  78. elif p.lower().endswith(".png"):
  79. yield p
  80. if __name__ == "__main__":
  81. n = 0
  82. for f in _walk(sys.argv[1:] or ["."]):
  83. try:
  84. if declutter_file(f):
  85. n += 1
  86. print("cleaned:", f)
  87. except Exception as e:
  88. print("skip", f, e)
  89. print(f"done. {n} file(s) cleaned.")