散称训练数据前置处理

data_preprocessing/extend.py

@@ -0,0 +1,131 @@
+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)