156 lines
6.0 KiB
Python
156 lines
6.0 KiB
Python
import numpy as np
|
||
from PIL import Image
|
||
from skimage import color
|
||
from scipy.optimize import minimize
|
||
import os
|
||
import time
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
|
||
# ================= 配置区域 =================
|
||
|
||
max_workers = 8 # 【并行】 并行处理的进程数 (建议设为 CPU 核心数，如 4, 8, 16)
|
||
# 图片后缀
|
||
prefix = "_AOD-Net" # "_10_0.2_result" # "_result"
|
||
# 输入文件夹 (带 prefix 后缀的图片)
|
||
src_dir = "AOD-Net"
|
||
# 参考文件夹 (GT/Ground Truth，无 prefix 后缀)
|
||
ref_dir = "去雾图像-北航合作-雾图"
|
||
# 输出文件夹
|
||
out_dir = "AOD-Net+后处理"
|
||
|
||
# ===========================================
|
||
|
||
# 尝试导入 tqdm，如果没安装则定义一个简单的占位符
|
||
try:
|
||
from tqdm import tqdm
|
||
except ImportError:
|
||
def tqdm(iterable, **kwargs):
|
||
return iterable
|
||
|
||
def process_single_image(filename, src_folder, ref_folder, output_folder, prefix=prefix):
|
||
"""
|
||
处理单张图片的函数，用于并行调用
|
||
"""
|
||
source_path = os.path.join(src_folder, filename)
|
||
|
||
try:
|
||
# --- 寻找对应的参考图 ---
|
||
# 逻辑：去除文件名后缀 "prefix"
|
||
name_no_ext, ext = os.path.splitext(filename)
|
||
|
||
if name_no_ext.endswith(prefix):
|
||
ref_name_no_ext = name_no_ext[:-len(prefix)] # 去掉最后n个字符 (prefix)
|
||
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"[跳过] 找不到参考图: {ref_filename} (对应: {filename})"
|
||
|
||
# --- 读取图片并归一化 ---
|
||
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
|
||
|
||
# --- 转换到 HSV ---
|
||
hsv_src = color.rgb2hsv(img_src)
|
||
hsv_ref = color.rgb2hsv(img_ref)
|
||
|
||
# --- 定义损失函数 ---
|
||
# 注意：在多进程中，loss_function 必须定义在 worker 内部才能访问到 hsv_src/ref
|
||
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
|
||
|
||
# --- 开始优化 ---
|
||
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)
|
||
|
||
# --- 应用最佳参数 ---
|
||
hsv_new = hsv_src.copy()
|
||
hsv_new[:, :, 1] = np.clip(hsv_new[:, :, 1] * best_s, 0, 1)
|
||
hsv_new[:, :, 2] = np.clip(hsv_new[:, :, 2] * best_v, 0, 1)
|
||
|
||
# --- 转回 RGB 并保存 ---
|
||
img_result_rgb = color.hsv2rgb(hsv_new)
|
||
img_save = Image.fromarray((img_result_rgb * 255).astype(np.uint8))
|
||
|
||
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)
|
||
|
||
return f"OK: {filename} -> S={s_percent}%, V={v_percent}%"
|
||
|
||
except Exception as e:
|
||
return f"[错误] 处理文件 {filename} 时出错: {str(e)}"
|
||
|
||
|
||
def calculate_and_process_batch_parallel(src_folder, ref_folder, output_folder, max_workers=None):
|
||
# 1. 确保输出目录存在
|
||
if not os.path.exists(output_folder):
|
||
os.makedirs(output_folder)
|
||
print(f"已创建输出目录: {output_folder}")
|
||
|
||
# 2. 获取源文件夹内所有图片文件
|
||
valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tif')
|
||
file_list = [f for f in os.listdir(src_folder) if f.lower().endswith(valid_extensions)]
|
||
|
||
total_files = len(file_list)
|
||
print(f"共发现 {total_files} 张图片，准备开始并行处理 (进程数: {max_workers if max_workers else '自动'})...\n")
|
||
|
||
# 3. 并行处理
|
||
results = []
|
||
|
||
# ProcessPoolExecutor 自动管理进程池
|
||
# max_workers=None 意味着使用 CPU 核心数
|
||
with ProcessPoolExecutor(max_workers=max_workers) as executor:
|
||
# 提交所有任务
|
||
future_to_file = {
|
||
executor.submit(process_single_image, filename, src_folder, ref_folder, output_folder): filename
|
||
for filename in file_list
|
||
}
|
||
|
||
# 使用 tqdm 显示进度条，as_completed 会在任何一个任务完成时yield
|
||
pbar = tqdm(total=total_files, unit="img")
|
||
|
||
for future in as_completed(future_to_file):
|
||
result_msg = future.result()
|
||
pbar.update(1)
|
||
|
||
# 如果是错误或跳过信息，打印出来；如果是OK，只更新进度条不刷屏（可选）
|
||
if not result_msg.startswith("OK"):
|
||
tqdm.write(result_msg) # 使用 tqdm.write 防止打断进度条
|
||
# else:
|
||
# tqdm.write(result_msg) # 如果想看每张图的详细结果，取消注释这行
|
||
|
||
pbar.close()
|
||
|
||
print("\n" + "="*30)
|
||
print("所有处理已完成。")
|
||
|
||
# 执行
|
||
if __name__ == "__main__":
|
||
# Windows 下使用多进程必须放在 if __name__ == "__main__": 之下
|
||
if os.path.exists(src_dir) and os.path.exists(ref_dir):
|
||
# max_workers 可以手动指定，例如 max_workers=4。如果不填则默认跑满 CPU。
|
||
calculate_and_process_batch_parallel(src_dir, ref_dir, out_dir, max_workers = max_workers)
|
||
else:
|
||
print("错误: 找不到输入文件夹或参考文件夹，请检查路径。") |