Python批量处理图片对比度并且保存

在别人帖子上进行了改进，加入了读取文件夹下的图片并进行批量处理并保存的模块

import cv2
import sys
import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
"""
  对比度调整算法：
  主要是对RGB空间进行调整。设定合适的RGB阈值，并在此阈值基础上计算出合适的调整系数进行对比度调整。
  参考CSDN博客：https://blog.csdn.net/maozefa/article/details/7069001
"""
def ContrastAlgorithm(rgb_img, contrast=0.5, threshold=0.5):
    img = rgb_img * 1.0
    img_out = img
    # 增量等于1，按灰度阈值最多调整成八种颜色：
    # 黑、红、绿、蓝、黄(255,255,0)、品红(255,0,255)、青(0,255,255)、白
    if contrast == 1:
        # newRGB = RGB >= Threshold? 255 : 0
        mask_1 = img >= threshold * 255.0
        rgb1 = 255.0
        rgb2 = 0
        img_out = rgb1 * mask_1 + rgb2 * (1 - mask_1)
    # 增量大于0小于1
    elif contrast >= 0:
        alpha = 1 - contrast
        alpha = 1 / alpha - 1
        img_out[:, :, 0] = img[:, :, 0] + (img[:, :, 0] - threshold * 255.0) * alpha
        img_out[:, :, 1] = img[:, :, 1] + (img[:, :, 1] - threshold * 255.0) * alpha
        img_out[:, :, 2] = img[:, :, 2] + (img[:, :, 2] - threshold * 255.0) * alpha
    # 增量小于0
    else:
        alpha = contrast
        img_out[:, :, 0] = img[:, :, 0] + (img[:, :, 0] - threshold * 255.0) * alpha
        img_out[:, :, 1] = img[:, :, 1] + (img[:, :, 1] - threshold * 255.0) * alpha
        img_out[:, :, 2] = img[:, :, 2] + (img[:, :, 2] - threshold * 255.0) * alpha
    # img_out = img_out / 255.0
    return img_out
#批量修改对比度并保存
def Save2dir(data_dir,save_dir):
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    for name in os.listdir(data_dir):
        save_name = name[:-4] + "hh.png"
        img = cv2.imread(os.path.join(data_dir,name))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_new = ContrastAlgorithm(img, contrast, threshold)
        # plt.imshow(img_new)
        # plt.savefig(os.path.join(save_dir,save_name))
        img_new = img_new[:, :, ::-1]
        cv2.imwrite(os.path.join(save_dir,save_name), img_new)


#  run : python Contrast.py (path) (contrast) (threshold)
def show():
    path = '/home/ljs/Desktop/test/Screenshot from bearing.mp4.png' #单张图片测试
    img = cv2.imread(path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_new = ContrastAlgorithm(img, contrast, threshold)
    plt.figure("img_original")
    plt.imshow(img / 255.0)
    plt.axis('off')

    plt.figure("img_contrast")
    plt.imshow(img_new / 255.0)
    plt.axis('off')
    plt.show()

def save():
    Data_dir = "/home/ljs/Desktop/test"  # 要改变的图片的路径文件夹
    Save_dir = "/home/ljs/Desktop/test"  # 要保存的图片的路径文件夹
    Save2dir(Data_dir,Save_dir)


if __name__ == '__main__':
    contrast = 0.5  # 范围：-1至1
    threshold = 0.5  # 范围：0至1
    show() # 查看一张图的效果
    # save() # 进行批量处理并且保存