# pip install opencv-python numpy
import os
import cv2
import numpy as np
from PIL import ImageGrab, Image

def adjust_gamma(image, gamma=1.0):
    # 建立一个查找表用于Gamma校正
    invGamma = 1.0 / gamma
    table = np.array([(i / 255.0) ** invGamma * 255
                      for i in np.arange(0, 256)]).astype("uint8")
    return cv2.LUT(image, table)

def enhance_document(input_image_path, output_image_path):
    print(f"正在处理 {input_image} ...")
    # 读取输入图像
    image = Image.open(input_image_path)
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # 转换为 OpenCV 格式

    # 检查图像是否加载成功
    if image is None:
        print(f"无法加载输入文件: {input_image_path}")
        return

    # 转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 使用自适应阈值进行二值化处理，去除背景
    binary = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 15, 15
    )

    # 使用形态学操作去除噪声，增强文字
    kernel = np.ones((1, 1), np.uint8)
    morph = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    # 进行Gamma校正，调整亮度和对比度
    enhanced = adjust_gamma(morph, gamma=1.5)

    # 保存输出图像
    pil_image = Image.fromarray(cv2.cvtColor(enhanced, cv2.COLOR_BGR2RGB))  # 转换回 PIL.Image
    pil_image.save(output_image_path)  # 使用 PIL 的 save 方法保存图像

    print(f"　　结果 {output_image_path}")


if __name__ == "__main__":
    # 获取当前工作目录下的所有图片文件
    supported_formats = ['.jpg', '.png', '.bmp']
    current_dir = os.getcwd()
    scan_dir = os.path.join(current_dir, "scan")
    if not os.path.exists(scan_dir):
        os.makedirs(scan_dir)
    
    # 遍历目录下的所有文件，筛选出支持的图片格式
    image_files = [
        file for file in os.listdir(current_dir) 
        if os.path.splitext(file)[1].lower() in supported_formats
    ]
    
    # 依次处理每个图片文件
    for image_file in image_files:
        input_image = os.path.join(current_dir, image_file)  # 输入的图片文件路径
        output_image = os.path.join(scan_dir, image_file)  # 输出的图片文件路径放到 'scan' 目录
        enhance_document(input_image, output_image)