训练数据前置处理与提升训练效率

data_preprocessing/extend.py

@@ -5,6 +5,9 @@ from PIL import Image, ImageEnhance
 
 
 class ImageExtendProcessor:
+    def __init__(self, conf):
+        self.conf = conf
+
     def is_image_file(self, filename):
         """检查文件是否为图像文件"""
         image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff')
@@ -80,7 +83,7 @@ class ImageExtendProcessor:
             print(f"处理图像 {image_path} 时出错: {e}")
             return False
 
-    def process_image_directory(self, src_dir, dst_dir, same_directory, **kwargs):
+    def process_extra_directory(self, src_dir, dst_dir, same_directory, dir_name):
         """
         处理单个目录中的图像文件
         :param src_dir: 源目录路径
@@ -99,19 +102,24 @@ class ImageExtendProcessor:
             if not same_directory:
                 # 复制原始文件  （另存文件夹时启用）
                 shutil.copy2(src_path, os.path.join(dst_dir, img_file))
+            if dir_name == 'extra':
+                # 生成并保存旋转后的图像
+                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)
+            elif dir_name == 'train':
+                # 生成并保存旋转后的图像
+                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 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):
+    def image_extend(self, src_dir, dst_dir, same_directory=False, dir_name=None):
         if same_directory:
             n_dst_dir = src_dir
             print(f"处理目录 {src_dir} 中的图像文件  保存至同一目录下")
@@ -121,7 +129,60 @@ class ImageExtendProcessor:
         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 dir_name == 'extra':
+                self.process_extra_directory(src_subdir_path,
+                                             dst_subdir_path,
+                                             same_directory,
+                                             dir_name)
+            if dir_name == 'train':
+                if len(os.listdir(src_subdir_path)) < 50:
+                    self.process_extra_directory(src_subdir_path,
+                                                 dst_subdir_path,
+                                                 same_directory,
+                                                 dir_name)
+
+    def random_remove_image(self, subdir_path, max_count=200):
+        """
+        随机删除子目录中的图像文件，直到数量不超过max_count
+        :param subdir_path: 子目录路径
+        :param max_count: 最大允许的图像数量
+        """
+        # 统计图像文件数量
+        image_files = [f for f in os.listdir(subdir_path)
+                       if self.is_image_file(f) and os.path.isfile(os.path.join(subdir_path, f))]
+        current_count = len(image_files)
+
+        # 如果图像数量不超过max_count，则无需删除
+        if current_count <= max_count:
+            print(f"无需处理 {subdir_path} (包含 {current_count} 个图像)")
+            return
+
+        # 计算需要删除的文件数
+        remove_count = current_count - max_count
+
+        # 随机选择要删除的文件
+        files_to_remove = random.sample(image_files, remove_count)
+
+        # 删除选中的文件
+        for file in files_to_remove:
+            file_path = os.path.join(subdir_path, file)
+            os.remove(file_path)
+            print(f"已删除 {file_path}")
+
+    def control_number(self):
+        if self.conf['extend']['extend_extra']:
+            self.image_extend(self.conf['extend']['extend_extra_dir'],
+                              '',
+                              same_directory=self.conf['extend']['extend_same_dir'],
+                              dir_name='extra')
+        if self.conf['extend']['extend_train']:
+            self.image_extend(self.conf['extend']['extend_train_dir'],
+                              '',
+                              same_directory=self.conf['extend']['extend_same_dir'],
+                              dir_name='train')
+        if self.conf['limit']['count_limit']:
+            self.random_remove_image(self.conf['limit']['limit_dir'],
+                                     max_count=self.conf['limit']['limit_count'])
 
 
 if __name__ == "__main__":