DAY 43 复习日

作业：

kaggle找到一个图像数据集，用cnn网络进行训练并且用grad-cam做可视化

进阶：并拆分成多个文件

import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import layers, models, applications
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
from tensorflow.keras.models import load_model
import cv2

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

class DataLoader:
    def __init__(self, data_dir, img_size=(224, 224), batch_size=32):
        self.data_dir = data_dir
        self.img_size = img_size
        self.batch_size = batch_size
        
    def load_data(self):
        """加载并预处理图像数据"""
        # 数据增强配置
        train_datagen = ImageDataGenerator(
            rescale=1./255,
            rotation_range=20,
            width_shift_range=0.2,
            height_shift_range=0.2,
            shear_range=0.2,
            zoom_range=0.2,
            horizontal_flip=True,
            validation_split=0.2
        )
        
        test_datagen = ImageDataGenerator(rescale=1./255)
        
        # 生成训练集和验证集
        train_generator = train_datagen.flow_from_directory(
            os.path.join(self.data_dir, 'train'),
            target_size=self.img_size,
            batch_size=self.batch_size,
            class_mode='categorical',
            subset='training'
        )
        
        val_generator = train_datagen.flow_from_directory(
            os.path.join(self.data_dir, 'train'),
            target_size=self.img_size,
            batch_size=self.batch_size,
            class_mode='categorical',
            subset='validation'
        )
        
        # 生成测试集
        test_generator = test_datagen.flow_from_directory(
            os.path.join(self.data_dir, 'test'),
            target_size=self.img_size,
            batch_size=self.batch_size,
            class_mode='categorical'
        )
        
        return train_generator, val_generator, test_generator

class CNNModel:
    def __init__(self, input_shape, num_classes):
        self.input_shape = input_shape
        self.num_classes = num_classes
        
    def build_simple_cnn(self):
        """构建简单的CNN模型"""
        model = models.Sequential([
            layers.Conv2D(32, (3, 3), activation='relu', input_shape=self.input_shape),
            layers.MaxPooling2D((2, 2)),
            layers.Conv2D(64, (3, 3), activation='relu'),
            layers.MaxPooling2D((2, 2)),
            layers.Conv2D(128, (3, 3), activation='relu'),
            layers.MaxPooling2D((2, 2)),
            layers.Flatten(),
            layers.Dense(128, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(self.num_classes, activation='softmax')
        ])
        
        model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        
        return model
    
    def build_pretrained_model(self, model_name='vgg16'):
        """构建基于预训练模型的CNN"""
        if model_name == 'vgg16':
            base_model = applications.VGG16(
                weights='imagenet',
                include_top=False,
                input_shape=self.input_shape
            )
        elif model_name == 'resnet50':
            base_model = applications.ResNet50(
                weights='imagenet',
                include_top=False,
                input_shape=self.input_shape
            )
        else:
            raise ValueError("不支持的预训练模型")
            
        # 冻结预训练层
        for layer in base_model.layers:
            layer.trainable = False
            
        # 添加自定义层
        x = base_model.output
        x = layers.GlobalAveragePooling2D()(x)
        x = layers.Dense(256, activation='relu')(x)
        x = layers.Dropout(0.5)(x)
        predictions = layers.Dense(self.num_classes, activation='softmax')(x)
        
        model = models.Model(inputs=base_model.input, outputs=predictions)
        
        model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        
        return model, base_model

class ModelTrainer:
    def __init__(self, model, model_path='best_model.h5'):
        self.model = model
        self.model_path = model_path
        
    def train(self, train_generator, val_generator, epochs=10):
        """训练模型"""
        callbacks = [
            EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=True),
            ModelCheckpoint(self.model_path, monitor='val_accuracy', save_best_only=True),
            ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=2, min_lr=0.00001)
        ]
        
        history = self.model.fit(
            train_generator,
            steps_per_epoch=train_generator.samples // train_generator.batch_size,
            validation_data=val_generator,
            validation_steps=val_generator.samples // val_generator.batch_size,
            epochs=epochs,
            callbacks=callbacks
        )
        
        return history

class GradCAM:
    def __init__(self, model, class_names, layer_name=None):
        self.model = model
        self.class_names = class_names
        
        # 如果没有指定层名，尝试自动找到最后一个卷积层
        if layer_name is None:
            layer_name = self._find_last_conv_layer()
        self.layer_name = layer_name
        
    def _find_last_conv_layer(self):
        """自动查找模型的最后一个卷积层"""
        for layer in reversed(self.model.layers):
            if 'conv' in layer.name:
                return layer.name
        raise ValueError("模型中没有找到卷积层")
        
    def generate_heatmap(self, img_array, pred_index=None):
        """生成Grad-CAM热力图"""
        # 创建一个用于获取输出的模型
        grad_model = tf.keras.models.Model(
            [self.model.inputs], 
            [self.model.get_layer(self.layer_name).output, self.model.output]
        )
        
        # 计算梯度
        with tf.GradientTape() as tape:
            conv_outputs, predictions = grad_model(img_array)
            if pred_index is None:
                pred_index = tf.argmax(predictions[0])
            class_channel = predictions[:, pred_index]
            
        # 获取梯度
        grads = tape.gradient(class_channel, conv_outputs)
        
        # 平均梯度
        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
        
        # 权重激活映射
        conv_outputs = conv_outputs[0]
        heatmap = tf.reduce_mean(tf.multiply(pooled_grads, conv_outputs), axis=-1)
        
        # 归一化热力图
        heatmap = np.maximum(heatmap, 0) / np.max(heatmap)
        return heatmap
    
    def overlay_heatmap(self, heatmap, img, alpha=0.4):
        """将热力图叠加到原图上"""
        # 调整热力图大小以匹配原图
        heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
        
        # 将热力图转换为RGB
        heatmap = np.uint8(255 * heatmap)
        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        
        # 将热力图叠加到原图上
        superimposed_img = heatmap * alpha + img
        superimposed_img = np.uint8(255 * superimposed_img / np.max(superimposed_img))
        
        return superimposed_img
    
    def visualize(self, img_path, img_size=(224, 224), alpha=0.4, top_n=3):
        """可视化Grad-CAM结果"""
        # 加载和预处理图像
        img = tf.keras.preprocessing.image.load_img(img_path, target_size=img_size)
        img_array = tf.keras.preprocessing.image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = img_array / 255.0
        
        # 预测类别
        predictions = self.model.predict(img_array)
        top_indices = np.argsort(predictions[0])[::-1][:top_n]
        
        # 生成热力图
        heatmap = self.generate_heatmap(img_array)
        
        # 加载原图用于显示
        original_img = cv2.imread(img_path)
        original_img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)
        original_img = cv2.resize(original_img, img_size)
        
        # 叠加热力图
        superimposed_img = self.overlay_heatmap(heatmap, original_img, alpha)
        
        # 显示结果
        plt.figure(figsize=(15, 5))
        
        plt.subplot(131)
        plt.title('原始图像')
        plt.imshow(original_img)
        plt.axis('off')
        
        plt.subplot(132)
        plt.title('Grad-CAM热力图')
        plt.imshow(heatmap, cmap='jet')
        plt.axis('off')
        
        plt.subplot(133)
        plt.title('叠加结果')
        plt.imshow(superimposed_img)
        plt.axis('off')
        
        # 显示预测结果
        plt.figtext(0.5, 0.01, 
                   f"预测结果:\n" + 
                   "\n".join([f"{self.class_names[idx]}: {predictions[0][idx]:.2%}" for idx in top_indices]),
                   ha="center", fontsize=12)
        
        plt.tight_layout()
        plt.show()

def download_kaggle_dataset(api_command):
    """使用Kaggle API下载数据集"""
    print(f"正在下载数据集: {api_command}")
    os.system(f"kaggle datasets download {api_command} -p ./data --unzip")
    print("数据集下载完成")

def main():
    # 设置参数
    data_dir = './data/your_dataset'  # 数据集路径
    img_size = (224, 224)
    batch_size = 32
    epochs = 10
    model_path = 'best_model.h5'
    use_pretrained = True  # 是否使用预训练模型
    
    # 下载Kaggle数据集（取消注释并提供正确的API命令）
    # download_kaggle_dataset('username/dataset-name')
    
    # 加载数据
    data_loader = DataLoader(data_dir, img_size, batch_size)
    train_generator, val_generator, test_generator = data_loader.load_data()
    num_classes = len(train_generator.class_indices)
    class_names = list(train_generator.class_indices.keys())
    
    # 构建模型
    cnn_model = CNNModel(input_shape=(*img_size, 3), num_classes=num_classes)
    
    if use_pretrained:
        model, base_model = cnn_model.build_pretrained_model()
        last_conv_layer = base_model.get_layer('block5_conv3').name  # VGG16最后一个卷积层
    else:
        model = cnn_model.build_simple_cnn()
        last_conv_layer = None  # 自动查找最后一个卷积层
    
    # 训练模型
    trainer = ModelTrainer(model, model_path)
    history = trainer.train(train_generator, val_generator, epochs)
    
    # 评估模型
    test_loss, test_acc = model.evaluate(test_generator)
    print(f"测试准确率: {test_acc:.2%}")
    
    # 可视化训练历史
    plt.figure(figsize=(12, 4))
    plt.subplot(121)
    plt.plot(history.history['accuracy'])
    plt.plot(history.history['val_accuracy'])
    plt.title('模型准确率')
    plt.ylabel('准确率')
    plt.xlabel('训练轮次')
    plt.legend(['训练', '验证'], loc='upper left')
    
    plt.subplot(122)
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('模型损失')
    plt.ylabel('损失')
    plt.xlabel('训练轮次')
    plt.legend(['训练', '验证'], loc='upper left')
    plt.tight_layout()
    plt.show()
    
    # 加载最佳模型用于Grad-CAM
    best_model = load_model(model_path)
    
    # 选择一个测试图像进行Grad-CAM可视化
    test_image_path = os.path.join(data_dir, 'test', class_names[0], os.listdir(os.path.join(data_dir, 'test', class_names[0]))[0])
    
    # 创建Grad-CAM对象并可视化
    grad_cam = GradCAM(best_model, class_names, last_conv_layer)
    grad_cam.visualize(test_image_path)

if __name__ == "__main__":
    main()