import os
import random
import shutil
from PIL import Image, ImageEnhance


class ImageExtendProcessor:
    def is_image_file(self, filename):
        """检查文件是否为图像文件"""
        image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff')
        return filename.lower().endswith(image_extensions)

    def random_cute_image(self, image_path, output_path, ratio=0.8):
        """
        对图像进行随机裁剪
        :param image_path: 输入图像路径
        :param output_path: 输出图像路径
        :param ratio: 裁剪比例，决定裁剪区域的大小，默认为0.8
        """
        try:
            with Image.open(image_path) as img:
                # 获取图像尺寸
                width, height = img.size

                # 计算裁剪后的尺寸
                new_width = int(width * ratio)
                new_height = int(height * ratio)

                # 随机生成裁剪起始点
                left = random.randint(0, width - new_width)
                top = random.randint(0, height - new_height)
                right = left + new_width
                bottom = top + new_height

                # 执行裁剪
                cropped_img = img.crop((left, top, right, bottom))

                # 保存裁剪后的图像
                cropped_img.save(output_path)
                return True
        except Exception as e:
            print(f"处理图像 {image_path} 时出错: {e}")
            return False

    def random_brightness(self, image_path, output_path, brightness_factor=None):
        """
        对图像进行随机亮度调整
        :param image_path: 输入图像路径
        :param output_path: 输出图像路径
        :param brightness_factor: 亮度调整因子，默认为随机值
        """
        try:
            with Image.open(image_path) as img:
                # 创建一个ImageEnhance.Brightness对象
                enhancer = ImageEnhance.Brightness(img)

                # 如果没有指定亮度因子，则随机生成
                if brightness_factor is None:
                    brightness_factor = random.uniform(0.5, 1.5)

                # 应用亮度调整
                brightened_img = enhancer.enhance(brightness_factor)

                # 保存调整后的图像
                brightened_img.save(output_path)
                return True
        except Exception as e:
            print(f"处理图像 {image_path} 时出错: {e}")
            return False

    def rotate_image(self, image_path, output_path, degrees):
        """旋转图像并保存到指定路径"""
        try:
            with Image.open(image_path) as img:
                # 旋转图像并自动调整画布大小
                rotated = img.rotate(degrees, expand=True)
                rotated.save(output_path)
                return True
        except Exception as e:
            print(f"处理图像 {image_path} 时出错: {e}")
            return False

    def process_image_directory(self, src_dir, dst_dir, same_directory, **kwargs):
        """
        处理单个目录中的图像文件
        :param src_dir: 源目录路径
        :param dst_dir: 目标目录路径
        """
        if not os.path.exists(dst_dir):
            os.makedirs(dst_dir)
        # 获取目录中所有图像文件
        image_files = [f for f in os.listdir(src_dir)
                       if self.is_image_file(f) and os.path.isfile(os.path.join(src_dir, f))]

        # 处理每个图像文件
        for img_file in image_files:
            src_path = os.path.join(src_dir, img_file)
            base_name, ext = os.path.splitext(img_file)
            if not same_directory:
                # 复制原始文件  （另存文件夹时启用）
                shutil.copy2(src_path, os.path.join(dst_dir, img_file))

            # 生成并保存旋转后的图像
            for angle in [90, 180, 270]:
                dst_path = os.path.join(dst_dir, f"{base_name}_rotated_{angle}{ext}")
                self.rotate_image(src_path, dst_path, angle)
            for ratio in [0.8, 0.85, 0.9]:
                dst_path = os.path.join(dst_dir, f"{base_name}_cute_{ratio}{ext}")
                self.random_cute_image(src_path, dst_path, ratio)
            for brightness_factor in [0.8, 0.9, 1.0]:
                dst_path = os.path.join(dst_dir, f"{base_name}_brightness_{brightness_factor}{ext}")
                self.random_brightness(src_path, dst_path, brightness_factor)

    def image_extend(self, src_dir, dst_dir, same_directory=False, **kwargs):
        if same_directory:
            n_dst_dir = src_dir
            print(f"处理目录 {src_dir} 中的图像文件  保存至同一目录下")
        else:
            n_dst_dir = dst_dir
            print(f"处理目录 {src_dir} 中的图像文件  保存至不同目录下")
        for src_subdir in os.listdir(src_dir):
            src_subdir_path = os.path.join(src_dir, src_subdir)
            dst_subdir_path = os.path.join(n_dst_dir, src_subdir)
            self.process_image_directory(src_subdir_path, dst_subdir_path, same_directory)


if __name__ == "__main__":
    src_dir = "./scatter_mini"
    dst_dir = "./scatter_add"
    image_ex = ImageExtendProcessor()
    image_ex.image_extend(src_dir, dst_dir, same_directory=False)