161 lines
6.3 KiB
Python
161 lines
6.3 KiB
Python
import numpy as np
|
||
from PIL import Image
|
||
from skimage import color, exposure
|
||
from scipy.optimize import minimize
|
||
import os
|
||
import time
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
|
||
# ================= 配置区域 =================
|
||
|
||
# 1. 输入与输出文件夹
|
||
SRC_DIR = "去雾图像-北航合作-Result_Baidu" # 待处理图片文件夹
|
||
REF_DIR = "去雾图像-北航合作" # 参考图(GT)文件夹
|
||
OUT_DIR = "去雾图像-北航合作-Result_Baidu_Own_V3" # 结果输出文件夹
|
||
|
||
# 2. 功能开关
|
||
ENABLE_HIST_MATCH = True # 【开关】 True: 开启直方图匹配; False: 关闭
|
||
MAX_WORKERS = 4 # 【并行】 并行处理的进程数 (建议设为 CPU 核心数，如 4, 8, 16)
|
||
|
||
# ===========================================
|
||
|
||
def process_single_image(file_info):
|
||
"""
|
||
单个图片处理函数 (用于并行调用)
|
||
file_info: (filename, src_dir, ref_dir, out_dir, enable_hist)
|
||
"""
|
||
filename, src_folder, ref_folder, output_folder, use_hist = file_info
|
||
|
||
source_path = os.path.join(src_folder, filename)
|
||
|
||
# 1. 寻找对应的参考图
|
||
# 逻辑：去除文件名后缀 "_result" (例如 "image01_result.png" -> "image01.png")
|
||
name_no_ext, ext = os.path.splitext(filename)
|
||
if name_no_ext.endswith("_result"):
|
||
ref_name_no_ext = name_no_ext[:-7] # 去掉 "_result"
|
||
else:
|
||
ref_name_no_ext = name_no_ext
|
||
|
||
ref_filename = ref_name_no_ext + ext
|
||
ref_path = os.path.join(ref_folder, ref_filename)
|
||
|
||
if not os.path.exists(ref_path):
|
||
return f"[跳过] 找不到参考图: {filename}"
|
||
|
||
try:
|
||
# 2. 读取图片并归一化 (0-1 float)
|
||
img_src_pil = Image.open(source_path).convert('RGB')
|
||
img_src = np.array(img_src_pil) / 255.0
|
||
|
||
img_ref_pil = Image.open(ref_path).convert('RGB')
|
||
if img_src_pil.size != img_ref_pil.size:
|
||
img_ref_pil = img_ref_pil.resize(img_src_pil.size, Image.BILINEAR)
|
||
img_ref = np.array(img_ref_pil) / 255.0
|
||
|
||
# 3. RGB -> HSV
|
||
hsv_src = color.rgb2hsv(img_src)
|
||
hsv_ref = color.rgb2hsv(img_ref)
|
||
|
||
# === 新增功能: 直方图匹配 (Histogram Matching) ===
|
||
if use_hist:
|
||
# 分离通道
|
||
s_h, s_s, s_v = hsv_src[:,:,0], hsv_src[:,:,1], hsv_src[:,:,2]
|
||
r_h, r_s, r_v = hsv_ref[:,:,0], hsv_ref[:,:,1], hsv_ref[:,:,2]
|
||
|
||
# 对 S 和 V 通道进行直方图匹配
|
||
# 这会将 src 的分布形状强行调整为 ref 的分布形状
|
||
matched_s = exposure.match_histograms(s_s, r_s)
|
||
matched_v = exposure.match_histograms(s_v, r_v)
|
||
|
||
# 更新 hsv_src，后续的 minimize 将在此基础上进一步微调系数
|
||
hsv_src = np.stack([s_h, matched_s, matched_v], axis=-1)
|
||
|
||
# 4. 优化 S/V 乘数因子
|
||
# 即使做了直方图匹配，我们依然计算一个最佳的整体缩放系数，以确保整体误差最小
|
||
def loss_function(params):
|
||
ks, kv = params
|
||
adj_s = np.clip(hsv_src[:,:,1] * ks, 0, 1)
|
||
adj_v = np.clip(hsv_src[:,:,2] * kv, 0, 1)
|
||
loss_s = np.mean((adj_s - hsv_ref[:,:,1])**2)
|
||
loss_v = np.mean((adj_v - hsv_ref[:,:,2])**2)
|
||
return loss_s + loss_v
|
||
|
||
# 初始猜测 [1.0, 1.0]
|
||
res = minimize(loss_function, [1.0, 1.0], method='Nelder-Mead', tol=1e-4)
|
||
best_s, best_v = res.x
|
||
|
||
s_percent = int(best_s * 100)
|
||
v_percent = int(best_v * 100)
|
||
|
||
# 5. 应用最终参数
|
||
hsv_final = hsv_src.copy()
|
||
hsv_final[:, :, 1] = np.clip(hsv_final[:, :, 1] * best_s, 0, 1)
|
||
hsv_final[:, :, 2] = np.clip(hsv_final[:, :, 2] * best_v, 0, 1)
|
||
|
||
# 6. 保存结果
|
||
img_result_rgb = color.hsv2rgb(hsv_final)
|
||
img_save = Image.fromarray((img_result_rgb * 255).astype(np.uint8))
|
||
|
||
# 命名增加标识，如果开启了直方图匹配，可以在文件名加个标记(可选)，
|
||
# 这里保持您要求的格式: 原文件名_S_XX_V_XX.png
|
||
new_filename = f"{name_no_ext}_S_{s_percent}_V_{v_percent}{ext}"
|
||
save_path = os.path.join(output_folder, new_filename)
|
||
|
||
img_save.save(save_path)
|
||
|
||
match_tag = "[HistMatch]" if use_hist else "[Raw]"
|
||
return f"{match_tag} 完成: {new_filename} (S={s_percent}%, V={v_percent}%)"
|
||
|
||
except Exception as e:
|
||
return f"[错误] {filename}: {str(e)}"
|
||
|
||
def main():
|
||
# 1. 检查文件夹
|
||
if not os.path.exists(SRC_DIR) or not os.path.exists(REF_DIR):
|
||
print("错误: 输入或参考文件夹不存在。")
|
||
return
|
||
|
||
if not os.path.exists(OUT_DIR):
|
||
os.makedirs(OUT_DIR)
|
||
|
||
# 2. 获取文件列表
|
||
valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tif')
|
||
file_list = [f for f in os.listdir(SRC_DIR) if f.lower().endswith(valid_extensions)]
|
||
total_files = len(file_list)
|
||
|
||
if total_files == 0:
|
||
print("源文件夹为空。")
|
||
return
|
||
|
||
print(f"=== 开始处理 ===")
|
||
print(f"模式: {'直方图匹配 + 参数优化' if ENABLE_HIST_MATCH else '仅参数优化'}")
|
||
print(f"并行: {MAX_WORKERS} 线程")
|
||
print(f"数量: {total_files} 张图片")
|
||
print("-" * 30)
|
||
|
||
# 3. 准备任务参数
|
||
tasks = []
|
||
for f in file_list:
|
||
# 打包参数传给 worker
|
||
tasks.append((f, SRC_DIR, REF_DIR, OUT_DIR, ENABLE_HIST_MATCH))
|
||
|
||
# 4. 并行执行
|
||
start_time = time.time()
|
||
|
||
with ProcessPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||
# 提交所有任务
|
||
futures = [executor.submit(process_single_image, task) for task in tasks]
|
||
|
||
# 获取结果 (as_completed 会在任务完成时立即返回)
|
||
for i, future in enumerate(as_completed(futures)):
|
||
result = future.result()
|
||
print(f"[{i+1}/{total_files}] {result}")
|
||
|
||
end_time = time.time()
|
||
print("-" * 30)
|
||
print(f"全部完成! 耗时: {end_time - start_time:.2f} 秒")
|
||
print(f"结果保存在: {OUT_DIR}")
|
||
|
||
if __name__ == "__main__":
|
||
# Windows 下使用多进程必须放在 if __name__ == "__main__": 之下
|
||
main() |