影像批量匀色（python）

from PIL import Image
import os
import numpy as np
import piexif

def histMatch(img, ref):
    out = np.zeros_like(img)
    _, _, colorChannel = img.shape
    for i in range(colorChannel):
        # get the histogram
        hist_img, _ = np.histogram(img[:, :, i], 256, range=(0, 255))
        hist_ref, _ = np.histogram(ref[:, :, i], 256, range=(0, 255))
        # get the accumulative histogram
        cdf_img = np.cumsum(hist_img)
        cdf_ref = np.cumsum(hist_ref)
        # match
        for j in range(1, 256):
            tmp = abs(cdf_img[j] - cdf_ref)
            tmp = tmp.tolist()
            # find the smallest number in tmp, get the index of this number
            idx = tmp.index(min(tmp))
            out[:, :, i][img[:, :, i] == j] = idx
    return out

def batchHistMatch(input_dir, ref_img_path, output_dir):
    # 批量处理指定文件夹下的图像，将它们与参考图像进行匀色，输出到指定文件夹
    ref_img = np.array(Image.open(ref_img_path).convert('RGB'))
    for file_name in os.listdir(input_dir):
        if file_name.endswith(".JPG") or file_name.endswith(".jpeg") or file_name.endswith(".png"):
            input_path = os.path.join(input_dir, file_name)
            output_path = os.path.join(output_dir, file_name)
            img = np.array(Image.open(input_path).convert('RGB'))
            matched_img = histMatch(img, ref_img)  # 调用上面实现的图像匹配函数
            exif_data = piexif.load(input_path)
            exif_bytes = piexif.dump(exif_data)
            Image.fromarray(matched_img.astype(np.uint8)).save(output_path, exif=exif_bytes)

# 指定输入文件夹路径、参考图像路径、输出文件夹路径
input_dir = ''
ref_img_path = ''
output_dir = ''
# 进行批量匀色处理
batchHistMatch(input_dir, ref_img_path, output_dir)