172 lines
5.6 KiB
Python
172 lines
5.6 KiB
Python
import cv2
|
||
import numpy as np
|
||
|
||
def detect_stain_mask(img2, img1, diff_thresh=0.12, grad_thresh=0.03, var_thresh=0.005):
|
||
"""
|
||
检测图2中的印染区域（相对于图1的差异 + 高梯度 + 高局部方差）
|
||
|
||
Args:
|
||
img2: 图2（清晰但有印染），uint8 RGB
|
||
img1: 图1（模糊但纯净），uint8 RGB
|
||
Returns:
|
||
mask: 浮点掩码 [0,1]，1表示印染区域
|
||
"""
|
||
# 转浮点
|
||
f2 = img2.astype(np.float32) / 255.0
|
||
f1 = img1.astype(np.float32) / 255.0
|
||
|
||
# 1. 差异图
|
||
diff = np.abs(f2 - f1)
|
||
gray_diff = cv2.cvtColor(diff, cv2.COLOR_RGB2GRAY)
|
||
|
||
# 2. 图2的拉普拉斯梯度（印染边缘高梯度）
|
||
gray2 = cv2.cvtColor(f2, cv2.COLOR_RGB2GRAY)
|
||
laplacian2 = cv2.Laplacian(gray2, cv2.CV_32F, ksize=3)
|
||
grad_mag = np.abs(laplacian2)
|
||
|
||
# 3. 图2的局部方差（印染区域局部对比度高）
|
||
mean2 = cv2.GaussianBlur(f2, (0, 0), 5)
|
||
sqr_mean2 = cv2.GaussianBlur(f2 * f2, (0, 0), 5)
|
||
local_var = np.mean(sqr_mean2 - mean2 * mean2, axis=2)
|
||
|
||
# 4. 组合条件
|
||
mask_cond = (
|
||
(gray_diff > diff_thresh) &
|
||
(grad_mag > grad_thresh) &
|
||
(local_var > var_thresh)
|
||
).astype(np.float32)
|
||
|
||
# 5. 形态学操作：膨胀连接区域 + 高斯模糊羽化
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
||
mask = cv2.dilate(mask_cond, kernel, iterations=2)
|
||
mask = cv2.GaussianBlur(mask, (15, 15), 0)
|
||
|
||
# 归一化到 [0, 1]
|
||
mask = np.clip(mask, 0, 1)
|
||
|
||
return mask
|
||
|
||
|
||
def fuse_textures(img2, img1, mask, sharpen_img1=True):
|
||
"""
|
||
融合两张纹理图
|
||
|
||
Args:
|
||
img2: 底图（清晰但有印染）
|
||
img1: 补丁图（模糊但纯净）
|
||
mask: 印染掩码 [0,1]
|
||
sharpen_img1: 是否对图1做轻微锐化以减少"糊"感
|
||
Returns:
|
||
fused: 融合结果
|
||
"""
|
||
f2 = img2.astype(np.float32) / 255.0
|
||
f1 = img1.astype(np.float32) / 255.0
|
||
|
||
# 可选：对图1做轻微锐化
|
||
if sharpen_img1:
|
||
# 非锐化掩模
|
||
blurred = cv2.GaussianBlur(f1, (0, 0), 1.5)
|
||
sharpened = cv2.addWeighted(f1, 1.5, blurred, -0.5, 0)
|
||
f1_sharp = np.clip(sharpened, 0, 1)
|
||
else:
|
||
f1_sharp = f1
|
||
|
||
# 融合：mask区域用图1，其余保留图2
|
||
mask_3ch = np.stack([mask, mask, mask], axis=2)
|
||
fused = f2 * (1.0 - mask_3ch) + f1_sharp * mask_3ch
|
||
|
||
# 转回 uint8
|
||
fused = np.clip(fused * 255.0, 0, 255).astype(np.uint8)
|
||
|
||
return fused
|
||
|
||
|
||
def post_process_clahe(img, clip_limit=1.5, grid_size=(8, 8), mix_ratio=0.3):
|
||
"""
|
||
轻微 CLAHE 后处理（仅 V 通道），与原图混合
|
||
"""
|
||
hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
|
||
h, s, v = cv2.split(hsv)
|
||
|
||
# CLAHE 需要 uint8
|
||
v_u8 = (v / 255.0 * 255).astype(np.uint8)
|
||
clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=grid_size)
|
||
v_adjusted = clahe.apply(v_u8).astype(np.float32)
|
||
|
||
# 混合
|
||
v_mixed = v * (1 - mix_ratio) + v_adjusted * mix_ratio
|
||
v_mixed = np.clip(v_mixed, 0, 255)
|
||
|
||
hsv_out = cv2.merge([h, s, v_mixed]).astype(np.uint8)
|
||
return cv2.cvtColor(hsv_out, cv2.COLOR_HSV2RGB)
|
||
|
||
|
||
def main():
|
||
# === 参数配置 ===
|
||
img1_path = "/home/algo/Documents/openMVS/openMVS/cartoon_texture/data/642713/642713.o.jpg" # 图1：模糊但纯净
|
||
img2_path = "/home/algo/Documents/openMVS/openMVS/cartoon_texture/data/642713/642713.mesh7.png" # 图2：清晰但有印染
|
||
output_path = "/home/algo/Documents/openMVS/openMVS/cartoon_texture/data/642713/642713.jpg"
|
||
|
||
# 印染检测参数（可按实际效果微调）
|
||
DIFF_THRESH = 0.12 # 差异阈值
|
||
GRAD_THRESH = 0.03 # 梯度阈值
|
||
VAR_THRESH = 0.005 # 局部方差阈值
|
||
|
||
# 后处理参数
|
||
APPLY_CLAHE = True # 是否做轻微CLAHE
|
||
SHARPEN_IMG1 = True # 是否锐化图1
|
||
|
||
# === 加载图片 ===
|
||
img2 = cv2.imread(img2_path)
|
||
img1 = cv2.imread(img1_path)
|
||
|
||
if img2 is None or img1 is None:
|
||
print("Error: Cannot load images")
|
||
return
|
||
|
||
# 确保尺寸一致（以图2为准）
|
||
if img1.shape != img2.shape:
|
||
print(f"Resizing img1 from {img1.shape} to {img2.shape}")
|
||
img1 = cv2.resize(img1, (img2.shape[1], img2.shape[0]))
|
||
|
||
# 转 RGB（OpenCV 默认 BGR）
|
||
img2_rgb = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
|
||
img1_rgb = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
|
||
|
||
print("Detecting stain mask...")
|
||
mask = detect_stain_mask(img2_rgb, img1_rgb,
|
||
diff_thresh=DIFF_THRESH,
|
||
grad_thresh=GRAD_THRESH,
|
||
var_thresh=VAR_THRESH)
|
||
|
||
print(f"Mask coverage: {np.mean(mask > 0.5) * 100:.1f}%")
|
||
|
||
print("Fusing textures...")
|
||
fused = fuse_textures(img2_rgb, img1_rgb, mask, sharpen_img1=SHARPEN_IMG1)
|
||
|
||
if APPLY_CLAHE:
|
||
print("Applying CLAHE post-processing...")
|
||
fused = post_process_clahe(fused)
|
||
|
||
# 保存结果
|
||
fused_bgr = cv2.cvtColor(fused, cv2.COLOR_RGB2BGR)
|
||
cv2.imwrite(output_path, fused_bgr)
|
||
|
||
# 保存掩码可视化（调试用）
|
||
mask_viz = (mask * 255).astype(np.uint8)
|
||
cv2.imwrite("stain_mask_debug.png", mask_viz)
|
||
|
||
print(f"Done! Result saved to {output_path}")
|
||
print(f"Debug mask saved to stain_mask_debug.png")
|
||
|
||
# 显示结果
|
||
cv2.imshow("Original (mesh7.1)", img2)
|
||
cv2.imshow("Reference (n1)", img1)
|
||
cv2.imshow("Stain Mask", mask_viz)
|
||
cv2.imshow("Fused Result", fused_bgr)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |