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


def calculate_clarity_score(image_path,target_size=(512, 512)):
    """
    计算图片的清晰度评分
    """
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    image = cv2.resize(image, target_size)
    if image is None:
        raise ValueError(f"无法读取图片：{image_path}")
    return cv2.Laplacian(image, cv2.CV_64F).var()


def process_images(input_folder, output_folder,clarity_score_low,clarity_score_high):
    if not os.path.exists(input_folder):
        print(f"输入文件夹不存在：{input_folder}")
        return

    for root, dirs, files in os.walk(input_folder):
        print(f'——————正在处理文件夹:{root}, 图片数量:{len(files)}——————')
        i = 0
        for filename in files:
            if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.jfif')):
                continue
            file_path = os.path.join(root, filename)
            # print(f'file_path is {file_path}')
            try:
                # 清晰度检测
                clarity_score = calculate_clarity_score(file_path)
                if clarity_score > clarity_score_low and clarity_score<clarity_score_high:
                    i+=1
                    # 读取图片
                    with Image.open(file_path) as img:
                        # resize
                        img = img.resize((512, 512))

                        # output_subfolder
                        relative_path = os.path.relpath(root, input_folder)
                        output_subfolder = os.path.join(output_folder, relative_path)
                            # print(f'root: {root}, input_folder: {input_folder}, relative_path: {relative_path}, output_subfolder: {output_subfolder}')                        
                        if not os.path.exists(output_subfolder):
                            os.makedirs(output_subfolder)
                        output_path = os.path.join(output_subfolder, os.path.splitext(filename)[0] + ".png")
                        img.save(output_path, format="PNG")
                    print(f"处理并保存图片{i}:{file_path} -> {output_path}********")
                else:
                    print(f"********图片 {file_path} 清晰度评分 {clarity_score} 不满足要求，跳过")
            except Exception as e:
                print(f"********处理图片 {file_path} 时出错：{e}")
        print(f'——————处理完毕,成功/图片数量: {i}/{len(files)}——————\n')


if __name__ == "__main__":
    input_folder = 'data_process/datasets/temp'  # 输入文件夹路径
    output_folder = 'data_process/datasets/target'  # 输出文件夹路径
    clarity_score_low = 50
    clarity_score_high =3000
    process_images(input_folder, output_folder,clarity_score_low,clarity_score_high)