YOLO视觉检测应用标注篇（四）

YOLO数据集图像预处理

准备自定义数据集时，首先要准备图片。
但收集到的图片可能规格各种各样。
也可能数量并不够。
因此我们需要进行预处理。
下面分享一个使用opencv进行预处理的程序：

import cv2
import os
import numpy as np
from PIL import Image, ExifTags
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm
import random

class YOLOPreprocessor:
    def __init__(self, config):
        self.config = {
            'target_size': 640,
            'use_denoise': True,
            'auto_denoise': True,
            'use_histogram': True,
            'color_enhance': True,
            'exif_correction': True,
            'data_augmentation': False,
            'normalize': True,
            'max_pixels': 3840*2160,
            'save_quality': 95,
            'num_workers': 4,
            'seed': 42
        }
        self.config.update(config)
        random.seed(self.config['seed'])
        np.random.seed(self.config['seed'])

    def _auto_denoise_check(self, gray_img):
        """自动判断是否需要去噪"""
        var = cv2.Laplacian(gray_img, cv2.CV_64F).var()
        return var < 50

    def _exif_orientation(self, img):
        """处理EXIF方向信息"""
        try:
            pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            exif = pil_img._getexif()
            if not exif:
                return img

            orientation_key = [k for k, v in ExifTags.TAGS.items() if v == 'Orientation']
            orientation = exif.get(orientation_key, 1)

            if orientation == 3:
                img = cv2.rotate(img, cv2.ROTATE_180)
            elif orientation == 6:
                img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
            elif orientation == 8:
                img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        except Exception as e:
            print(f"EXIF处理错误: {str(e)}")
        return img

    def _color_enhancement(self, img):
        """LAB颜色空间增强"""
        lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        return cv2.cvtColor(cv2.merge((l,a,b)), cv2.COLOR_LAB2RGB)

    def _data_augmentation(self, img):
        """随机数据增强"""
        # 随机水平翻转
        if random.random() > 0.5:
            img = cv2.flip(img, 1)
        
        # 随机旋转 (-15°到15°)
        angle = random.uniform(-15, 15)
        h, w = img.shape[:2]
        M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)
        img = cv2.warpAffine(img, M, (w, h), borderMode=cv2.BORDER_REPLICATE)
        
        # 随机亮度调整
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
        hsv[:,:,2] = hsv[:,:,2] * random.uniform(0.8, 1.2)
        img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
        
        return img

    def preprocess(self, img_path):
        """预处理流水线"""
        try:
            # 读取图像
            img = cv2.imread(img_path)
            if img is None:
                raise ValueError(f"无法读取图像: {img_path}")

            # EXIF方向校正
            if self.config['exif_correction']:
                img = self._exif_orientation(img)

            # 自动去噪判断
            if self.config['auto_denoise']:
                gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                if self._auto_denoise_check(gray):
                    img = cv2.GaussianBlur(img, (3,3), 0)

            # 手动去噪
            elif self.config['use_denoise']:
                img = cv2.GaussianBlur(img, (3,3), 0)

            # 转换为RGB
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            # 颜色增强
            if self.config['color_enhance']:
                img = self._color_enhancement(img)

            # 直方图均衡化
            if self.config['use_histogram']:
                yuv = cv2.cvtColor(img, cv2.COLOR_RGB2YUV)
                clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
                yuv[:,:,0] = clahe.apply(yuv[:,:,0])
                img = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB)

            # 数据增强
            if self.config['data_augmentation']:
                img = self._data_augmentation(img)

            # 尺寸处理
            h, w = img.shape[:2]
            
            # 内存保护
            if h * w > self.config['max_pixels']:
                scale = (self.config['max_pixels'] / (h * w)) ** 0.5
                h, w = int(h * scale), int(w * scale)
                img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)

            # 保持长宽比的缩放
            target_size = self.config['target_size']
            scale = min(target_size / h, target_size / w)
            new_h, new_w = int(h * scale), int(w * scale)
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

            # 填充
            padded = np.full((target_size, target_size, 3), 114, dtype=np.uint8)
            padded[:new_h, :new_w] = img

            # 归一化
            if self.config['normalize']:
                padded = padded.astype(np.float32) / 255.0

            return padded
        except Exception as e:
            print(f"处理 {os.path.basename(img_path)} 时出错: {str(e)}")
            return None

    def process_folder(self, input_dir, output_dir, save_format='.jpg'):
        """批量处理文件夹"""
        os.makedirs(output_dir, exist_ok=True)
        files = [f for f in os.listdir(input_dir) if f.lower().endswith('.bmp')]
        
        with ThreadPoolExecutor(max_workers=self.config['num_workers']) as executor:
            futures = []
            for filename in files:
                in_path = os.path.join(input_dir, filename)
                base_name = os.path.splitext(filename)[0] 
                out_filename = f"{base_name}{save_format}"
                out_path = os.path.join(output_dir, out_filename)
                futures.append(executor.submit(self._process_single, in_path, out_path))
            
            for future in tqdm(futures, total=len(files), desc="处理进度"):
                future.result()

    def _process_single(self, in_path, out_path):
        processed = self.preprocess(in_path)
        if processed is not None:
            if self.config['normalize']:
                save_img = (processed * 255).astype(np.uint8)
            else:
                save_img = processed
                
            save_img = cv2.cvtColor(save_img, cv2.COLOR_RGB2BGR)
            
            if out_path.lower().endswith('.jpg'):
                cv2.imwrite(out_path, save_img, [int(cv2.IMWRITE_JPEG_QUALITY), self.config['save_quality']])
            else:
                cv2.imwrite(out_path, save_img)


if __name__ == "__main__":
    # 配置参数（可根据需要修改）
    config = {
        'target_size': 640,
        'use_denoise': True,            # 手动去噪判断
        'auto_denoise': True,           # 自动去噪判断
        'use_histogram': False,         # 直方图均衡化
        'color_enhance': False,         # 颜色增强
        'exif_correction': False,       # EXIF方向校正
        'data_augmentation': False,     # 启用数据增强
        'normalize': True,              # 保存时需要归一化
        'num_workers': 2,               # 根据CPU核心数调整
        'save_quality': 95              # JPEG质量
    }

    # 初始化预处理器
    preprocessor = YOLOPreprocessor(config)
    
    # 执行处理
    preprocessor.process_folder(
        input_dir=r'path\to\pic', # 输入文件夹路径
        output_dir=r'path\to\output', # 输出文件夹路径
        save_format='.jpg'  # 或 '.bmp'
    )