data_preprocessing.py

import os
import numpy as np
import pandas as pd
from PIL import Image
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from sklearn.model_selection import KFold
import matplotlib.pyplot as plt
import seaborn as sns

class DataPreprocessor:
    """数据预处理类，负责图像数据的加载、预处理和数据增强"""
    
    def __init__(self, train_dir='train', img_size=(128, 128), batch_size=128):
        self.train_dir = train_dir
        self.img_size = img_size
        self.batch_size = batch_size
        self.class_names = ['cat', 'dog']
        
    def load_data_info(self):
        """加载数据集信息"""
        files = os.listdir(self.train_dir)
        cats = [f for f in files if f.startswith('cat')]
        dogs = [f for f in files if f.startswith('dog')]
        
        print(f"数据集统计:")
        print(f"猫的图片数量: {len(cats)}")
        print(f"狗的图片数量: {len(dogs)}")
        print(f"总图片数量: {len(files)}")
        
        return len(cats), len(dogs), len(files)
    
    def create_dataframe(self):
        """创建包含文件路径和标签的DataFrame"""
        files = os.listdir(self.train_dir)
        data = []
        
        for file in files:
            if file.startswith('cat'):
                label = 0  # cat
            elif file.startswith('dog'):
                label = 1  # dog
            else:
                continue
                
            data.append({
                'filename': file,
                'filepath': os.path.join(self.train_dir, file),
                'label': label,
                'class_name': 'cat' if label == 0 else 'dog'
            })
        
        df = pd.DataFrame(data)
        return df
    
    def create_data_generators(self, train_df, val_df):
        """创建训练和验证数据生成器 - 简化数据增强"""
        # 简化训练数据增强，减少计算量
        train_datagen = ImageDataGenerator(
            rescale=1./255,
            rotation_range=15,  # 减少旋转范围
            horizontal_flip=True,
            zoom_range=0.1,     # 减少缩放范围
            fill_mode='nearest'
        )
        
        # 验证数据只进行标准化
        val_datagen = ImageDataGenerator(rescale=1./255)
        
        # 创建训练生成器
        train_generator = train_datagen.flow_from_dataframe(
            train_df,
            x_col='filepath',
            y_col='class_name',
            target_size=self.img_size,
            batch_size=self.batch_size,
            class_mode='binary',
            shuffle=True
        )
        
        # 创建验证生成器
        val_generator = val_datagen.flow_from_dataframe(
            val_df,
            x_col='filepath',
            y_col='class_name',
            target_size=self.img_size,
            batch_size=self.batch_size,
            class_mode='binary',
            shuffle=False
        )
        
        return train_generator, val_generator
    
    def split_data_kfold(self, df, n_splits=3, random_state=42):
        """使用K折交叉验证分割数据 - 减少折数以加快实验"""
        kfold = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
        
        folds = []
        for fold_idx, (train_idx, val_idx) in enumerate(kfold.split(df)):
            train_df = df.iloc[train_idx].reset_index(drop=True)
            val_df = df.iloc[val_idx].reset_index(drop=True)
            
            folds.append({
                'fold': fold_idx + 1,
                'train_df': train_df,
                'val_df': val_df,
                'train_size': len(train_df),
                'val_size': len(val_df)
            })
            
            print(f"Fold {fold_idx + 1}: 训练集 {len(train_df)} 张, 验证集 {len(val_df)} 张")
        
        return folds
    
    def visualize_data_distribution(self, df):
        """可视化数据分布"""
        plt.figure(figsize=(12, 4))
        
        # 类别分布
        plt.subplot(1, 2, 1)
        class_counts = df['class_name'].value_counts()
        plt.pie(class_counts.values, labels=class_counts.index, autopct='%1.1f%%')
        plt.title('类别分布')
        
        # 类别计数
        plt.subplot(1, 2, 2)
        sns.countplot(data=df, x='class_name')
        plt.title('各类别图片数量')
        plt.ylabel('数量')
        
        plt.tight_layout()
        plt.savefig('data_distribution.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def sample_images(self, df, n_samples=8):
        """展示样本图片"""
        fig, axes = plt.subplots(2, 4, figsize=(16, 8))
        axes = axes.ravel()
        
        # 每个类别采样一半
        cat_samples = df[df['class_name'] == 'cat'].sample(n_samples//2)
        dog_samples = df[df['class_name'] == 'dog'].sample(n_samples//2)
        samples = pd.concat([cat_samples, dog_samples])
        
        for idx, (_, row) in enumerate(samples.iterrows()):
            img = Image.open(row['filepath'])
            img = img.resize(self.img_size)
            
            axes[idx].imshow(img)
            axes[idx].set_title(f"{row['class_name']}")
            axes[idx].axis('off')
        
        plt.tight_layout()
        plt.savefig('sample_images.png', dpi=300, bbox_inches='tight')
        plt.show()

if __name__ == "__main__":
    # 测试数据预处理功能
    preprocessor = DataPreprocessor()
    
    # 加载数据信息
    preprocessor.load_data_info()
    
    # 创建数据框
    df = preprocessor.create_dataframe()
    print(f"\n数据框形状: {df.shape}")
    print(f"数据框列名: {df.columns.tolist()}")
    print(f"\n前5行数据:")
    print(df.head())
    
    # 可视化数据分布
    preprocessor.visualize_data_distribution(df)
    
    # 展示样本图片
    preprocessor.sample_images(df)
    
    # K折交叉验证分割
    folds = preprocessor.split_data_kfold(df, n_splits=3)
    print(f"\n成功创建 {len(folds)} 个交叉验证折")