style_protector.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from tqdm import tqdm
import numpy as np
from models.simple_diffusion import SimpleDiffusionModel
from utils.noise_scheduler import NoiseScheduler

class ImageStyleProtector:
    """图像风格保护器"""
    
    def __init__(self, config, device='cuda', protection_strength=0.1, num_steps=50):
        self.device = device
        self.config = config
        self.protection_strength = protection_strength
        self.num_steps = num_steps
        
        # 初始化扩散模型
        self.diffusion_model = SimpleDiffusionModel(
            channels=config['model']['channels'],
            image_size=config['model']['image_size'],
            time_dim=config['model']['time_dim']
        ).to(device)
        
        # 初始化噪声调度器
        self.noise_scheduler = NoiseScheduler(
            num_timesteps=config['model']['num_timesteps'],
            device=device
        )
        
        # 确保所有组件都在正确的设备上
        self.diffusion_model = self.diffusion_model.to(device)
        self.noise_scheduler = self.noise_scheduler.to(device)
        
        # 加载分类器
        self.classifier = self._load_classifier()
        
        print(f"ImageStyleProtector 初始化完成")
        print(f"保护强度: {self.protection_strength}")
        print(f"扩散步数: {self.num_steps}")
    
    def _load_classifier(self):
        """加载预训练的CIFAR-10分类器"""
        model = models.resnet18(pretrained=True)
        model.fc = nn.Linear(model.fc.in_features, 10)  # CIFAR-10有10个类别
        model = model.to(self.device)
        model.eval()
        return model
    
    def protect_image(self, image):
        """保护单张图像"""
        with torch.no_grad():
            # 添加噪声
            noise = torch.randn_like(image) * self.protection_strength
            
            # 使用扩散过程进行保护
            protected = self._apply_diffusion_protection(image, noise)
            
            return protected
    
    def protect_batch(self, images):
        """批量保护图像"""
        protected_images = []
        
        for i in range(images.size(0)):
            protected = self.protect_image(images[i:i+1])
            protected_images.append(protected)
        
        return torch.cat(protected_images, dim=0)
    
    def _apply_diffusion_protection(self, image, noise):
        """应用扩散保护"""
        # 简化的扩散保护过程
        t = torch.randint(0, self.num_steps, (image.size(0),), device=self.device)
        
        # 添加噪声
        noisy_image = self.noise_scheduler.add_noise(image, noise, t)
        
        # 去噪（保护过程）
        protected = self._denoise_step(noisy_image, t)
        
        return protected
    
    def _denoise_step(self, noisy_image, t):
        """去噪步骤"""
        # 简化的去噪过程
        with torch.no_grad():
            # 预测噪声
            predicted_noise = self.diffusion_model(noisy_image, t)
            
            # 去除噪声
            alpha = self.noise_scheduler.get_alpha(t)
            protected = (noisy_image - predicted_noise * (1 - alpha)) / alpha.sqrt()
            
            # 限制像素值范围
            protected = torch.clamp(protected, -1, 1)
            
            return protected
    
    def train_model(self, train_loader, epochs, lr):
        """训练保护模型"""
        optimizer = torch.optim.Adam(self.diffusion_model.parameters(), lr=lr)
        criterion = nn.MSELoss()
        
        self.diffusion_model.train()
        
        for epoch in range(epochs):
            total_loss = 0
            num_batches = 0
            
            pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{epochs}')
            
            for images, _ in pbar:
                images = images.to(self.device)
                
                # 随机时间步
                t = torch.randint(0, self.num_steps, (images.size(0),), device=self.device)
                
                # 添加噪声
                noise = torch.randn_like(images)
                noisy_images = self.noise_scheduler.add_noise(images, noise, t)
                
                # 预测噪声
                predicted_noise = self.diffusion_model(noisy_images, t)
                
                # 计算损失
                loss = criterion(predicted_noise, noise)
                
                # 反向传播
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                total_loss += loss.item()
                num_batches += 1
                
                pbar.set_postfix({'Loss': f'{loss.item():.4f}'})
            
            avg_loss = total_loss / num_batches
            print(f'Epoch {epoch+1}/{epochs}, Average Loss: {avg_loss:.4f}')
    
    def evaluate_protection(self, test_loader):
        """评估保护效果"""
        self.diffusion_model.eval()
        self.classifier.eval()
        
        total_protection = 0
        total_quality = 0
        original_correct = 0
        protected_correct = 0
        total_samples = 0
        
        with torch.no_grad():
            for images, labels in tqdm(test_loader, desc='评估保护效果'):
                images = images.to(self.device)
                labels = labels.to(self.device)
                
                # 保护图像
                protected_images = self.protect_batch(images)
                
                # 计算分类准确率
                original_pred = self.classifier(images)
                protected_pred = self.classifier(protected_images)
                
                original_correct += (original_pred.argmax(1) == labels).sum().item()
                protected_correct += (protected_pred.argmax(1) == labels).sum().item()
                
                # 计算保护强度（分类置信度下降）
                original_conf = F.softmax(original_pred, dim=1).max(1)[0]
                protected_conf = F.softmax(protected_pred, dim=1).max(1)[0]
                protection = (original_conf - protected_conf).mean().item()
                total_protection += protection
                
                # 计算图像质量（MSE）
                quality = 1 - F.mse_loss(images, protected_images).item()
                total_quality += quality
                
                total_samples += images.size(0)
                
                if len(test_loader) > 10 and total_samples > 500:  # 限制评估样本数
                    break
        
        num_batches = min(len(test_loader), 10)
        
        results = {
            'avg_protection': total_protection / num_batches,
            'image_quality': total_quality / num_batches,
            'original_accuracy': original_correct / total_samples,
            'protected_accuracy': protected_correct / total_samples,
            'accuracy_drop': (original_correct - protected_correct) / total_samples
        }
        
        return results
    
    def save_model(self, path):
        """保存模型"""
        torch.save({
            'diffusion_model': self.diffusion_model.state_dict(),
            'config': self.config,
            'protection_strength': self.protection_strength,
            'num_steps': self.num_steps
        }, path)
    
    def load_model(self, path):
        """加载模型"""
        checkpoint = torch.load(path, map_location=self.device)
        self.diffusion_model.load_state_dict(checkpoint['diffusion_model'])
        self.config = checkpoint['config']
        self.protection_strength = checkpoint['protection_strength']
        self.num_steps = checkpoint['num_steps']