对三个不同文件夹中的图像进行相似度匹配，并将匹配成功的三张图像（每个文件夹各一张）拼接成一张横向长图保存。以下是详细的功能解析：

核心功能

1. 图像匹配： • 从三个文件夹（yuantu原图、Effect_ox1效果图1、Effect_pf效果图2）中读取图片。

   • 通过计算图像的均方误差（MSE）衡量相似度，找到每个原图对应的最相似的两张效果图（分别来自两个效果图文件夹）。

   • 使用全局优化算法（迭代移除已匹配的图片）确保匹配的唯一性和最优性。

2. 图像拼接： • 将匹配成功的三张图片（原图 + 效果图1 + 效果图2）横向拼接成一张长图。

   • 拼接时保持所有图片高度一致，宽度按原始比例缩放。

3. 结果保存： • 将拼接后的图像保存到save_mse文件夹中，文件名按序号命名（如001.jpg）。

关键步骤

1. 预处理图像： • 将所有图片统一缩放到384x384像素（用于MSE计算），确保尺寸一致。

   • 使用PIL.Image.LANCZOS保持缩放质量。

2. 相似度计算： • 计算原图与两个效果图文件夹中每张图片的MSE距离（值越小越相似）。

   • 构建距离矩阵（原图 vs 效果图1，原图 vs 效果图2）。

3. 优化匹配： • 通过迭代选择当前剩余图片中的最佳匹配（MSE总和最小），避免重复匹配。

   • 每次匹配后从候选池中移除已匹配的图片。

4. 拼接逻辑： • 使用原始尺寸图片拼接（非缩放后的384x384图片）。

   • 保持原图高度不变，调整效果图宽度以匹配原图高度（保持宽高比）。

输入输出 • 输入：

• 三个文件夹的图片路径：

◦ `heji/yuantu/`：原图

◦ `heji/Effect_ox1/`：效果图1

◦ `heji/Effect_pf/`：效果图2

• 支持格式：.png, .jpg, .jpeg

• 输出：

• 拼接后的长图保存在save_mse/目录下，按序号命名（如001.jpg）。

技术细节

1. MSE计算： • 公式：MSE = mean((img1 - img2)^2)，数值越小表示越相似。

   • 处理异常：如果图片无法打开或尺寸不匹配，返回无限大距离（float('inf')）。

2. 进度显示： • 使用tqdm库显示进度条（若未安装则回退到简单打印）。

3. 错误处理： • 捕获图片打开、缩放、拼接时的异常，跳过错误文件。

python
import os
import numpy as np
from PIL import Image
import io
from pathlib import Path
from collections import defaultdict
import time
from scipy.optimize import linear_sum_assignment
try:
    from tqdm import tqdm
    tqdm_available = True
except ImportError:
    tqdm_available = False
    print("提示: 安装 tqdm 库可以获得更好的进度条显示 (pip install tqdm)")

# Helper function to resize images for MSE calculation
def resize_for_mse(img_path, target_size=384):
    """Resizes an image to target_size x target_size for MSE comparison."""
    try:
        img = Image.open(img_path)
        img = img.convert('RGB')  # Ensure RGB mode for consistency
        img = img.resize((target_size, target_size), Image.LANCZOS)
        return np.array(img)
    except Exception as e:
        print(f"Error resizing image {os.path.basename(img_path)}: {e}")
        return None

# New function to just open images without resizing for final collage
def open_original_image(img_path):
    """Opens an image without resizing, for full resolution collages."""
    try:
        return Image.open(img_path)
    except Exception as e:
        print(f"Error opening image {os.path.basename(img_path)}: {e}")
        return None

# Calculate MSE between two image arrays
def calculate_mse(img1_array, img2_array):
    """Calculate Mean Squared Error between two image arrays."""
    if img1_array is None or img2_array is None:
        return float('inf')  # Max distance if one image is missing
    
    # Ensure arrays have the same shape
    if img1_array.shape != img2_array.shape:
        print(f"Warning: Image arrays have different shapes. Cannot compute MSE.")
        return float('inf')
    
    # Calculate MSE
    mse = np.mean((img1_array.astype(float) - img2_array.astype(float)) ** 2)
    return mse

# 创建保存目录
save_dir = "save_mse"
os.makedirs(save_dir, exist_ok=True)

# 获取所有图片路径
yuantu_dir = os.path.join("heji", "yuantu")
Effect_ox1_dir = os.path.join("heji", "Effect_ox1")
Effect_pf_dir = os.path.join("heji", "Effect_pf")

yuantu_images = [os.path.join(yuantu_dir, f) for f in os.listdir(yuantu_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
Effect_ox1_images = [os.path.join(Effect_ox1_dir, f) for f in os.listdir(Effect_ox1_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
Effect_pf_images = [os.path.join(Effect_pf_dir, f) for f in os.listdir(Effect_pf_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

# --- 预处理图像 ---
print("正在预处理图像 (调整大小到 384x384)...")

# 预处理所有图像
yuantu_arrays = {}
Effect_ox1_arrays = {}
Effect_pf_arrays = {}

# 处理原图
print("处理原图...")
yuantu_iter = tqdm(yuantu_images) if tqdm_available else yuantu_images
for image_path in yuantu_iter:
    if not tqdm_available:
        print(f"处理 {len(yuantu_arrays)+1}/{len(yuantu_images)}: {os.path.basename(image_path)}", end="\r")
    
    img_array = resize_for_mse(image_path)
    if img_array is not None:
        yuantu_arrays[image_path] = img_array

if not tqdm_available:
    print()
print(f"原图处理完成: {len(yuantu_arrays)}/{len(yuantu_images)}")

# 处理Effect_ox1
print("处理Effect_ox1...")
Effect_ox1_iter = tqdm(Effect_ox1_images) if tqdm_available else Effect_ox1_images
for image_path in Effect_ox1_iter:
    if not tqdm_available:
        print(f"处理 {len(Effect_ox1_arrays)+1}/{len(Effect_ox1_images)}: {os.path.basename(image_path)}", end="\r")
    
    img_array = resize_for_mse(image_path)
    if img_array is not None:
        Effect_ox1_arrays[image_path] = img_array

if not tqdm_available:
    print()
print(f"Effect_ox1处理完成: {len(Effect_ox1_arrays)}/{len(Effect_ox1_images)}")

# 处理Effect_pf
print("处理Effect_pf...")
Effect_pf_iter = tqdm(Effect_pf_images) if tqdm_available else Effect_pf_images
for image_path in Effect_pf_iter:
    if not tqdm_available:
        print(f"处理 {len(Effect_pf_arrays)+1}/{len(Effect_pf_images)}: {os.path.basename(image_path)}", end="\r")
    
    img_array = resize_for_mse(image_path)
    if img_array is not None:
        Effect_pf_arrays[image_path] = img_array

if not tqdm_available:
    print()
print(f"Effect_pf处理完成: {len(Effect_pf_arrays)}/{len(Effect_pf_images)}")

print("图像预处理完成！")

# --- 全局优化匹配算法 ---
print("开始全局优化匹配...")

# 将所有图片转换为索引列表，方便后续处理
valid_yuantu_images = list(yuantu_arrays.keys())
valid_Effect_ox1_images = list(Effect_ox1_arrays.keys())
valid_Effect_pf_images = list(Effect_pf_arrays.keys())

print(f"有效图片数量: 原图={len(valid_yuantu_images)}, Effect_ox1={len(valid_Effect_ox1_images)}, Effect_pf={len(valid_Effect_pf_images)}")

if len(valid_yuantu_images) == 0 or len(valid_Effect_ox1_images) == 0 or len(valid_Effect_pf_images) == 0:
    print("错误: 至少一个文件夹中没有有效的图片，无法进行匹配。")
    exit(1)

# 创建距离矩阵：yuantu-Effect_ox1和yuantu-Effect_pf
print("计算所有图片对之间的MSE距离...")

# 计算yuantu与Effect_ox1之间的距离矩阵
yuantu_Effect_ox1_distances = np.zeros((len(valid_yuantu_images), len(valid_Effect_ox1_images)))
for i, yuantu_path in enumerate(valid_yuantu_images):
    yuantu_arr = yuantu_arrays[yuantu_path]
    for j, Effect_ox1_path in enumerate(valid_Effect_ox1_images):
        Effect_ox1_arr = Effect_ox1_arrays[Effect_ox1_path]
        yuantu_Effect_ox1_distances[i, j] = calculate_mse(yuantu_arr, Effect_ox1_arr)

# 计算yuantu与Effect_pf之间的距离矩阵
yuantu_Effect_pf_distances = np.zeros((len(valid_yuantu_images), len(valid_Effect_pf_images)))
for i, yuantu_path in enumerate(valid_yuantu_images):
    yuantu_arr = yuantu_arrays[yuantu_path]
    for j, Effect_pf_path in enumerate(valid_Effect_pf_images):
        Effect_pf_arr = Effect_pf_arrays[Effect_pf_path]
        yuantu_Effect_pf_distances[i, j] = calculate_mse(yuantu_arr, Effect_pf_arr)

print("MSE距离计算完成，开始优化匹配...")

# 优化匹配逻辑
matched_triplets = []
remaining_yuantu = list(range(len(valid_yuantu_images)))
remaining_Effect_ox1 = list(range(len(valid_Effect_ox1_images)))
remaining_Effect_pf = list(range(len(valid_Effect_pf_images)))

# 因为我们需要三个文件夹的最优匹配，我们会迭代地移除已匹配的图片
iteration = 0
max_iterations = min(len(valid_yuantu_images), len(valid_Effect_ox1_images), len(valid_Effect_pf_images))

while (iteration < max_iterations and 
       len(remaining_yuantu) > 0 and 
       len(remaining_Effect_ox1) > 0 and 
       len(remaining_Effect_pf) > 0):
    
    iteration += 1
    print(f"匹配迭代 {iteration}/{max_iterations}，剩余: 原图={len(remaining_yuantu)}, Effect_ox1={len(remaining_Effect_ox1)}, Effect_pf={len(remaining_Effect_pf)}")
    
    # 构建当前迭代的子距离矩阵
    curr_yuantu_Effect_ox1 = np.zeros((len(remaining_yuantu), len(remaining_Effect_ox1)))
    for i, yuantu_idx in enumerate(remaining_yuantu):
        for j, Effect_ox1_idx in enumerate(remaining_Effect_ox1):
            curr_yuantu_Effect_ox1[i, j] = yuantu_Effect_ox1_distances[yuantu_idx, Effect_ox1_idx]
    
    curr_yuantu_Effect_pf = np.zeros((len(remaining_yuantu), len(remaining_Effect_pf)))
    for i, yuantu_idx in enumerate(remaining_yuantu):
        for j, Effect_pf_idx in enumerate(remaining_Effect_pf):
            curr_yuantu_Effect_pf[i, j] = yuantu_Effect_pf_distances[yuantu_idx, Effect_pf_idx]
    
    # 合并距离矩阵，寻找最佳组合
    best_total_distance = float('inf')
    best_triplet = None
    
    for i in range(len(remaining_yuantu)):
        yuantu_idx = remaining_yuantu[i]
        
        # 找到与当前yuantu最相似的Effect_ox1
        Effect_ox1_distances = curr_yuantu_Effect_ox1[i, :]
        min_Effect_ox1_distance = np.min(Effect_ox1_distances)
        min_Effect_ox1_j = np.argmin(Effect_ox1_distances)
        Effect_ox1_idx = remaining_Effect_ox1[min_Effect_ox1_j]
        
        # 找到与当前yuantu最相似的Effect_pf
        Effect_pf_distances = curr_yuantu_Effect_pf[i, :]
        min_Effect_pf_distance = np.min(Effect_pf_distances)
        min_Effect_pf_j = np.argmin(Effect_pf_distances)
        Effect_pf_idx = remaining_Effect_pf[min_Effect_pf_j]
        
        # 计算总距离
        total_distance = min_Effect_ox1_distance + min_Effect_pf_distance
        
        # 更新最佳匹配
        if total_distance < best_total_distance:
            best_total_distance = total_distance
            best_triplet = (yuantu_idx, Effect_ox1_idx, Effect_pf_idx, min_Effect_ox1_j, min_Effect_pf_j)
    
    if best_triplet:
        yuantu_idx, Effect_ox1_idx, Effect_pf_idx, Effect_ox1_local_idx, Effect_pf_local_idx = best_triplet
        matched_triplets.append((valid_yuantu_images[yuantu_idx], 
                                valid_Effect_ox1_images[Effect_ox1_idx], 
                                valid_Effect_pf_images[Effect_pf_idx]))
        
        # 从可用索引中移除已匹配的图片
        remaining_yuantu.remove(yuantu_idx)
        remaining_Effect_ox1.remove(Effect_ox1_idx)
        remaining_Effect_pf.remove(Effect_pf_idx)
    else:
        # 如果无法找到最佳匹配，退出循环
        print("无法找到更多的最佳匹配，结束匹配过程。")
        break

print(f"成功匹配了 {len(matched_triplets)} 组三元组。")

# --- 拼接和保存图像 ---
print("开始拼接和保存图像...")
print("使用原始尺寸图像进行拼接，可能会生成较大文件...")
matched_iter = tqdm(enumerate(matched_triplets), total=len(matched_triplets)) if tqdm_available else enumerate(matched_triplets)

for i, (yuantu_path, Effect_ox1_path, Effect_pf_path) in matched_iter:
    if not tqdm_available:
        print(f"拼接图片 {i+1}/{len(matched_triplets)}", end="\r")
    
    try:
        # 打开原始尺寸图片，不进行预缩放
        yuantu_img = open_original_image(yuantu_path)
        Effect_ox1_img = open_original_image(Effect_ox1_path)
        Effect_pf_img = open_original_image(Effect_pf_path)
        
        if not yuantu_img or not Effect_ox1_img or not Effect_pf_img:
            print(f"无法打开三元组中的一个或多个图像，跳过: {os.path.basename(yuantu_path)}")
            continue
        
        # 获取原始尺寸
        yuantu_w, yuantu_h = yuantu_img.size
        Effect_ox1_w, Effect_ox1_h = Effect_ox1_img.size
        Effect_pf_w, Effect_pf_h = Effect_pf_img.size
        
        # 计算最大高度，使用原图的高度作为参考
        target_height = yuantu_h
        
        # 调整其他图片高度与原图一致，保持宽高比
        new_Effect_ox1_w = int(Effect_ox1_w * target_height / Effect_ox1_h)
        Effect_ox1_img = Effect_ox1_img.resize((new_Effect_ox1_w, target_height), Image.LANCZOS)
        
        new_Effect_pf_w = int(Effect_pf_w * target_height / Effect_pf_h)
        Effect_pf_img = Effect_pf_img.resize((new_Effect_pf_w, target_height), Image.LANCZOS)
        
        # 创建新图像
        total_width = yuantu_w + new_Effect_ox1_w + new_Effect_pf_w
        new_img = Image.new('RGB', (total_width, target_height))
        
        # 粘贴图像
        current_x = 0
        new_img.paste(yuantu_img, (current_x, 0))
        current_x += yuantu_w
        new_img.paste(Effect_ox1_img, (current_x, 0))
        current_x += new_Effect_ox1_w
        new_img.paste(Effect_pf_img, (current_x, 0))
        
        # 保存拼接图像
        save_path = os.path.join(save_dir, f"{i+1:03d}.jpg")
        new_img.save(save_path, quality=95)  # 使用高质量保存
        
    except Exception as e:
        print(f"拼接图片时出错 ({i+1}): {e}")
        continue

print("\n所有图片拼接完成！共生成 {len(matched_triplets)} 张拼接图像。")