109_机器学习隐写技术深度解析：从GAN到深度学习的智能数据隐藏实现指南

输入: 载体图像 + 秘密信息 → 编码器 → 隐写图像 → 解码器 → 提取的秘密信息

生成器: 载体图像 + 秘密信息 → 隐写图像
判别器: 判断输入图像是原始载体还是隐写载体

编码器: 载体 + 秘密信息 → 隐写图像
解码器: 隐写图像 → 提取的秘密信息

min_G max_D V(D, G) = E_{x~P_data(x)}[log D(x)] + E_{z~P_z(z)}[log(1-D(G(z)))]

载体图像 + 秘密信息 → 隐写生成器 → 隐写图像
                    ↓
隐写判别器 ← 原始载体图像

# 创建并激活虚拟环境
python -m venv steganography_env
source steganography_env/bin/activate  # Linux/Mac
steganography_env\Scripts\activate  # Windows

# 安装核心依赖
pip install tensorflow-gpu==2.8.0  # 或 tensorflow==2.8.0（无GPU）
pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install numpy==1.21.0
pip install matplotlib==3.4.2
pip install pillow==8.3.1
pip install opencv-python==4.5.3
pip install scikit-learn==0.24.2
pip install tqdm==4.61.2
pip install jupyter==1.0.0

import tensorflow as tf
import torch
import numpy as np
import matplotlib.pyplot as plt
import cv2
import sklearn

# 验证TensorFlow和CUDA
tf_version = tf.__version__
cuda_available = tf.test.is_built_with_cuda()
gpu_available = len(tf.config.list_physical_devices('GPU')) > 0
print(f"TensorFlow版本: {tf_version}")
print(f"CUDA支持: {cuda_available}")
print(f"GPU可用: {gpu_available}")

# 验证PyTorch和CUDA
torch_version = torch.__version__
torch_cuda_available = torch.cuda.is_available()
if torch_cuda_available:
    print(f"PyTorch版本: {torch_version}")
    print(f"CUDA可用: {torch_cuda_available}")
    print(f"CUDA设备数: {torch.cuda.device_count()}")
    print(f"当前CUDA设备: {torch.cuda.current_device()}")
    print(f"CUDA设备名称: {torch.cuda.get_device_name(0)}")

# 验证其他库
print(f"NumPy版本: {np.__version__}")
print(f"Matplotlib版本: {plt.matplotlib.__version__}")
print(f"OpenCV版本: {cv2.__version__}")
print(f"scikit-learn版本: {sklearn.__version__}")

秘密信息 →
           ↓
载体图像 → 嵌入网络 → 隐写图像 → 提取网络 → 重建的秘密信息
                ↓
原始载体图像 → 判别网络 → 分类结果

def build_encoder(input_shape, secret_shape):
    # 载体图像输入
    cover_input = Input(shape=input_shape, name='cover_input')
    # 秘密信息输入
    secret_input = Input(shape=secret_shape, name='secret_input')
    
    # 扩展秘密信息维度，使其与载体图像匹配
    secret_expanded = Conv2D(3, (1, 1), padding='same')(secret_input)
    secret_expanded = UpSampling2D(size=(input_shape[0]//secret_shape[0], input_shape[1]//secret_shape[1]))(secret_expanded)
    
    # 合并载体和秘密信息
    merged = Concatenate(axis=3)([cover_input, secret_expanded])
    
    # 编码器主体
    x = Conv2D(64, (3, 3), padding='same', activation='relu')(merged)
    x = Conv2D(64, (3, 3), padding='same', activation='relu')(x)
    x = Conv2D(3, (3, 3), padding='same', activation='tanh')(x)
    
    # 残差连接
    stego_output = Add()([cover_input, x])
    
    encoder = Model(inputs=[cover_input, secret_input], outputs=stego_output, name='encoder')
    return encoder

def build_decoder(input_shape, secret_shape):
    # 隐写图像输入
    stego_input = Input(shape=input_shape, name='stego_input')
    
    # 解码器主体
    x = Conv2D(64, (3, 3), padding='same', activation='relu')(stego_input)
    x = MaxPooling2D((2, 2))(x)
    x = Conv2D(128, (3, 3), padding='same', activation='relu')(x)
    x = MaxPooling2D((2, 2))(x)
    
    # 全连接层提取秘密信息
    x = Flatten()(x)
    x = Dense(np.prod(secret_shape), activation='sigmoid')(x)
    secret_output = Reshape(secret_shape)(x)
    
    decoder = Model(inputs=stego_input, outputs=secret_output, name='decoder')
    return decoder

def build_discriminator(input_shape):
    # 图像输入
    image_input = Input(shape=input_shape, name='image_input')
    
    # 判别器主体
    x = Conv2D(64, (3, 3), padding='same', strides=(2, 2), activation='leaky_relu')(image_input)
    x = Conv2D(128, (3, 3), padding='same', strides=(2, 2), activation='leaky_relu')(x)
    x = Conv2D(256, (3, 3), padding='same', strides=(2, 2), activation='leaky_relu')(x)
    
    # 分类层
    x = Flatten()(x)
    x = Dense(1, activation='sigmoid')(x)
    
    discriminator = Model(inputs=image_input, outputs=x, name='discriminator')
    return discriminator

L_recon = E[||s - D(E(c, s))||²]

L_cover = E[||c - E(c, s)||²]

L_adv = E[log(1 - D_dis(E(c, s)))]

L_total = α * L_recon + β * L_cover + γ * L_adv

# 下载CelebA数据集示例（使用gdown工具）
pip install gdown
gdown https://drive.google.com/uc?id=0B7EVK8r0v71pZjFTYXZWM3FlRnM -O img_align_celeba.zip
unzip img_align_celeba.zip -d ./dataset/celeba/

# 使用PyTorch的torchvision下载CIFAR-10数据集
import torchvision
import torchvision.transforms as transforms

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transforms.ToTensor())
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())

# 将图像归一化到[-1, 1]范围
def normalize_image(image):
    return (image / 127.5) - 1.0

# 将图像归一化到[0, 1]范围
def normalize_image_01(image):
    return image / 255.0

import cv2
import numpy as np

def preprocess_image(image_path, target_size=(256, 256)):
    # 读取图像
    image = cv2.imread(image_path)
    # 转换为RGB格式
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 调整图像尺寸
    image = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA)
    # 转换为numpy数组
    image = np.array(image, dtype=np.float32)
    # 归一化
    image = normalize_image(image)
    return image

from PIL import Image
from torchvision import transforms

def get_image_transforms():
    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomRotation(10),
        transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
    return transform

# 计算数据集的均值和标准差
def calculate_mean_std(dataset):
    mean = np.zeros(3)
    std = np.zeros(3)
    total_samples = 0
    
    for image in dataset:
        image = np.array(image) / 255.0
        mean += np.mean(image, axis=(0, 1))
        std += np.std(image, axis=(0, 1))
        total_samples += 1
    
    mean /= total_samples
    std /= total_samples
    return mean, std

# 使用计算出的均值和标准差标准化图像
def standardize_image(image, mean, std):
    return (image / 255.0 - mean) / std

import librosa

def convert_audio_format(input_path, output_path):
    # 加载音频
    y, sr = librosa.load(input_path, sr=None)
    # 保存为WAV格式
    librosa.output.write_wav(output_path, y, sr)

def resample_audio(input_path, target_sr=16000):
    y, sr = librosa.load(input_path, sr=target_sr)
    return y, target_sr

def segment_audio(audio, sr, segment_length=3):
    # 将音频分割成指定长度的片段
    segment_samples = int(segment_length * sr)
    segments = []
    
    for i in range(0, len(audio), segment_samples):
        segment = audio[i:i+segment_samples]
        # 如果最后一个片段不足指定长度，则进行零填充
        if len(segment) < segment_samples:
            segment = np.pad(segment, (0, segment_samples - len(segment)), 'constant')
        segments.append(segment)
    
    return np.array(segments)

def extract_melspectrogram(audio, sr, n_mels=128, n_fft=2048, hop_length=512):
    # 计算梅尔频谱图
    melspectrogram = librosa.feature.melspectrogram(
        y=audio, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length
    )
    # 转换为对数刻度
    melspectrogram = librosa.power_to_db(melspectrogram, ref=np.max)
    return melspectrogram

def text_to_binary(text):
    # 将文本转换为二进制字符串
    binary = ''.join(format(ord(char), '08b') for char in text)
    return binary

def binary_to_text(binary):
    # 将二进制字符串转换为文本
    text = ''
    # 确保二进制字符串的长度是8的倍数
    binary = binary.zfill((len(binary) + 7) // 8 * 8)
    for i in range(0, len(binary), 8):
        byte = binary[i:i+8]
        text += chr(int(byte, 2))
    return text

def pad_data(data, target_length):
    # 计算需要填充的数据量
    pad_length = target_length - len(data)
    if pad_length > 0:
        # 使用零进行填充
        data += '0' * pad_length
    return data

def split_data(data, block_size):
    # 将数据分成指定大小的块
    blocks = []
    for i in range(0, len(data), block_size):
        blocks.append(data[i:i+block_size])
    return blocks

def preprocess_secret_image(image_path, target_size=(64, 64)):
    # 读取并预处理秘密图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # 转换为灰度图像以减少数据量
    image = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA)
    image = np.array(image, dtype=np.float32) / 255.0  # 归一化到[0, 1]范围
    return image

import zlib

def compress_data(data):
    # 压缩数据
    if isinstance(data, str):
        data_bytes = data.encode('utf-8')
    else:
        data_bytes = data
    compressed = zlib.compress(data_bytes)
    return compressed

def decompress_data(compressed_data):
    # 解压缩数据
    decompressed = zlib.decompress(compressed_data)
    return decompressed

from cryptography.fernet import Fernet

def generate_key():
    # 生成加密密钥
    key = Fernet.generate_key()
    return key

def encrypt_data(data, key):
    # 加密数据
    f = Fernet(key)
    if isinstance(data, str):
        data_bytes = data.encode('utf-8')
    else:
        data_bytes = data
    encrypted = f.encrypt(data_bytes)
    return encrypted

def decrypt_data(encrypted_data, key):
    # 解密数据
    f = Fernet(key)
    decrypted = f.decrypt(encrypted_data)
    return decrypted

from torchvision import transforms

def get_image_augmentation():
    augmentation = transforms.Compose([
        transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.2),
        transforms.RandomRotation(degrees=15),
        transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1)
    ])
    return augmentation

def add_gaussian_noise(image, mean=0, std=0.01):
    # 添加高斯噪声
    noise = np.random.normal(mean, std, image.shape)
    noisy_image = image + noise
    # 裁剪到有效范围
    noisy_image = np.clip(noisy_image, 0, 1)
    return noisy_image

def add_salt_and_pepper_noise(image, salt_prob=0.01, pepper_prob=0.01):
    # 添加椒盐噪声
    noisy_image = np.copy(image)
    # 添加盐噪声（白色像素）
    num_salt = int(np.ceil(salt_prob * image.size))
    coords = [np.random.randint(0, i - 1, int(num_salt)) for i in image.shape]
    noisy_image[tuple(coords)] = 1
    # 添加胡椒噪声（黑色像素）
    num_pepper = int(np.ceil(pepper_prob * image.size))
    coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in image.shape]
    noisy_image[tuple(coords)] = 0
    return noisy_image

def apply_blur(image, kernel_size=(3, 3)):
    # 应用高斯模糊
    blurred_image = cv2.GaussianBlur(image, kernel_size, 0)
    return blurred_image

def adjust_volume(audio, volume_factor=1.5):
    # 调整音量
    adjusted_audio = audio * volume_factor
    # 裁剪到有效范围
    adjusted_audio = np.clip(adjusted_audio, -1, 1)
    return adjusted_audio

def add_audio_noise(audio, noise_level=0.01):
    # 添加高斯噪声
    noise = np.random.normal(0, noise_level, len(audio))
    noisy_audio = audio + noise
    # 裁剪到有效范围
    noisy_audio = np.clip(noisy_audio, -1, 1)
    return noisy_audio

def time_stretch(audio, stretch_factor=1.1):
    # 时间拉伸（不改变音高）
    stretched_audio = librosa.effects.time_stretch(audio, rate=stretch_factor)
    return stretched_audio

def pitch_shift(audio, sr, n_steps=2):
    # 音高变换（不改变速度）
    shifted_audio = librosa.effects.pitch_shift(audio, sr=sr, n_steps=n_steps)
    return shifted_audio

class SteganographyDataset(torch.utils.data.Dataset):
    def __init__(self, cover_images, secret_images, transform=None):
        self.cover_images = cover_images
        self.secret_images = secret_images
        self.transform = transform
    
    def __len__(self):
        return len(self.cover_images)
    
    def __getitem__(self, idx):
        cover = self.cover_images[idx]
        secret = self.secret_images[idx]
        
        # 确保增强的一致性
        if self.transform:
            # 设置随机种子以确保相同的增强变换
            seed = np.random.randint(2147483647)
            
            # 对载体图像应用增强
            np.random.seed(seed)
            cover = self.transform(cover)
            
            # 对秘密图像应用增强
            np.random.seed(seed)
            secret = self.transform(secret)
        
        return {'cover': cover, 'secret': secret}

# 检查并配置CUDA环境
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import os

# 检查CUDA是否可用
print(f"CUDA可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU型号: {torch.cuda.get_device_name(0)}")
    device = torch.device("cuda:0")
else:
    print("使用CPU进行训练")
    device = torch.device("cpu")

# 设置随机种子以确保结果可复现
def set_seed(seed=42):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed()

# 定义数据加载器
def get_data_loaders(train_dataset, val_dataset, batch_size=32, num_workers=4):
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    return train_loader, val_loader

# 假设我们已经定义好了数据集
train_dataset = SteganographyDataset(train_cover_images, train_secret_images, transform=train_transform)
val_dataset = SteganographyDataset(val_cover_images, val_secret_images, transform=val_transform)

train_loader, val_loader = get_data_loaders(train_dataset, val_dataset, batch_size=32)

# 假设我们已经定义好了编码器、解码器和判别器
from models import Encoder, Decoder, Discriminator

# 初始化模型
encoder = Encoder(in_channels=3, out_channels=3).to(device)
decoder = Decoder(in_channels=3, out_channels=3).to(device)
discriminator = Discriminator(in_channels=3).to(device)

# 初始化优化器
optimizer_encoder = optim.Adam(encoder.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_decoder = optim.Adam(decoder.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_discriminator = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

# 定义损失函数
criterion_mse = nn.MSELoss().to(device)
criterion_bce = nn.BCEWithLogitsLoss().to(device)

def train_one_epoch(encoder, decoder, discriminator, train_loader, 
                   optimizer_encoder, optimizer_decoder, optimizer_discriminator,
                   criterion_mse, criterion_bce, device, alpha=1.0, beta=1.0, gamma=1.0):
    # 设置模型为训练模式
    encoder.train()
    decoder.train()
    discriminator.train()
    
    running_loss_encoder = 0.0
    running_loss_decoder = 0.0
    running_loss_discriminator = 0.0
    
    for i, data in enumerate(train_loader):
        # 获取数据
        cover_images = data['cover'].to(device)
        secret_images = data['secret'].to(device)
        
        # 标签
        batch_size = cover_images.size(0)
        real_labels = torch.ones(batch_size, 1).to(device)
        fake_labels = torch.zeros(batch_size, 1).to(device)
        
        # ============== 训练判别器 ==============
        # 判别真实图像
        real_outputs = discriminator(cover_images)
        d_loss_real = criterion_bce(real_outputs, real_labels)
        
        # 生成隐写图像
        stego_images = encoder(cover_images, secret_images)
        # 判别隐写图像
        fake_outputs = discriminator(stego_images.detach())
        d_loss_fake = criterion_bce(fake_outputs, fake_labels)
        
        # 总判别器损失
        d_loss = (d_loss_real + d_loss_fake) / 2
        
        # 反向传播和优化
        optimizer_discriminator.zero_grad()
        d_loss.backward()
        optimizer_discriminator.step()
        
        # ============== 训练编码器和解码器 ==============
        # 生成隐写图像
        stego_images = encoder(cover_images, secret_images)
        # 解码获取秘密图像
        decoded_secret = decoder(stego_images)
        
        # 重建损失
        reconstruction_loss = criterion_mse(decoded_secret, secret_images)
        
        # 不可感知性损失
        imperceptibility_loss = criterion_mse(stego_images, cover_images)
        
        # 对抗损失
        fake_outputs = discriminator(stego_images)
        adversarial_loss = criterion_bce(fake_outputs, real_labels)
        
        # 总损失
        e_loss = alpha * reconstruction_loss + beta * imperceptibility_loss + gamma * adversarial_loss
        d_loss_encoder = alpha * reconstruction_loss + beta * imperceptibility_loss + gamma * adversarial_loss
        
        # 反向传播和优化
        optimizer_encoder.zero_grad()
        optimizer_decoder.zero_grad()
        e_loss.backward()
        optimizer_encoder.step()
        optimizer_decoder.step()
        
        # 更新运行损失
        running_loss_encoder += e_loss.item() * batch_size
        running_loss_decoder += d_loss_encoder.item() * batch_size
        running_loss_discriminator += d_loss.item() * batch_size
    
    # 计算平均损失
    epoch_loss_encoder = running_loss_encoder / len(train_loader.dataset)
    epoch_loss_decoder = running_loss_decoder / len(train_loader.dataset)
    epoch_loss_discriminator = running_loss_discriminator / len(train_loader.dataset)
    
    return epoch_loss_encoder, epoch_loss_decoder, epoch_loss_discriminator

def validate(encoder, decoder, val_loader, criterion_mse, device):
    # 设置模型为评估模式
    encoder.eval()
    decoder.eval()
    
    running_reconstruction_loss = 0.0
    running_imperceptibility_loss = 0.0
    
    with torch.no_grad():
        for data in val_loader:
            cover_images = data['cover'].to(device)
            secret_images = data['secret'].to(device)
            
            # 生成隐写图像和解码秘密图像
            stego_images = encoder(cover_images, secret_images)
            decoded_secret = decoder(stego_images)
            
            # 计算损失
            reconstruction_loss = criterion_mse(decoded_secret, secret_images)
            imperceptibility_loss = criterion_mse(stego_images, cover_images)
            
            # 更新运行损失
            batch_size = cover_images.size(0)
            running_reconstruction_loss += reconstruction_loss.item() * batch_size
            running_imperceptibility_loss += imperceptibility_loss.item() * batch_size
    
    # 计算平均损失
    val_reconstruction_loss = running_reconstruction_loss / len(val_loader.dataset)
    val_imperceptibility_loss = running_imperceptibility_loss / len(val_loader.dataset)
    
    return val_reconstruction_loss, val_imperceptibility_loss

# 开始训练
epochs = 100
best_val_loss = float('inf')

# 用于记录训练过程
loss_history = {
    'encoder': [],
    'decoder': [],
    'discriminator': [],
    'val_reconstruction': [],
    'val_imperceptibility': []
}

for epoch in range(epochs):
    # 训练一个epoch
    train_e_loss, train_d_loss, train_disc_loss = train_one_epoch(
        encoder, decoder, discriminator, train_loader,
        optimizer_encoder, optimizer_decoder, optimizer_discriminator,
        criterion_mse, criterion_bce, device
    )
    
    # 验证
    val_rec_loss, val_imp_loss = validate(
        encoder, decoder, val_loader, criterion_mse, device
    )
    
    # 记录损失
    loss_history['encoder'].append(train_e_loss)
    loss_history['decoder'].append(train_d_loss)
    loss_history['discriminator'].append(train_disc_loss)
    loss_history['val_reconstruction'].append(val_rec_loss)
    loss_history['val_imperceptibility'].append(val_imp_loss)
    
    # 打印训练信息
    print(f'Epoch [{epoch+1}/{epochs}], '
          f'Train E Loss: {train_e_loss:.4f}, '
          f'Train D Loss: {train_d_loss:.4f}, '
          f'Train Disc Loss: {train_disc_loss:.4f}, '
          f'Val Rec Loss: {val_rec_loss:.4f}, '
          f'Val Imp Loss: {val_imp_loss:.4f}')
    
    # 保存最佳模型
    current_val_loss = val_rec_loss + val_imp_loss
    if current_val_loss < best_val_loss:
        best_val_loss = current_val_loss
        # 保存模型
        torch.save(encoder.state_dict(), 'best_encoder.pth')
        torch.save(decoder.state_dict(), 'best_decoder.pth')
        torch.save(discriminator.state_dict(), 'best_discriminator.pth')
        print(f'Saved best model at epoch {epoch+1}')

# 保存模型
def save_model(model, optimizer, epoch, loss, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss
    }, path)

# 加载模型
def load_model(model, optimizer, path):
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    loss = checkpoint['loss']
    return model, optimizer, epoch, loss

# 保存最终模型
save_model(encoder, optimizer_encoder, epochs, loss_history['encoder'][-1], 'final_encoder.pth')
save_model(decoder, optimizer_decoder, epochs, loss_history['decoder'][-1], 'final_decoder.pth')
save_model(discriminator, optimizer_discriminator, epochs, loss_history['discriminator'][-1], 'final_discriminator.pth')

# 使用学习率调度器
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR

# 1. StepLR: 每N个epoch将学习率乘以gamma
scheduler_encoder = StepLR(optimizer_encoder, step_size=20, gamma=0.5)
scheduler_decoder = StepLR(optimizer_decoder, step_size=20, gamma=0.5)
scheduler_discriminator = StepLR(optimizer_discriminator, step_size=20, gamma=0.5)

# 2. ReduceLROnPlateau: 当验证损失停止改善时降低学习率
scheduler_encoder = ReduceLROnPlateau(optimizer_encoder, mode='min', factor=0.5, patience=10, verbose=True)
scheduler_decoder = ReduceLROnPlateau(optimizer_decoder, mode='min', factor=0.5, patience=10, verbose=True)

# 3. CosineAnnealingLR: 使用余弦退火调度学习率
scheduler_encoder = CosineAnnealingLR(optimizer_encoder, T_max=epochs, eta_min=0.00001)
scheduler_decoder = CosineAnnealingLR(optimizer_decoder, T_max=epochs, eta_min=0.00001)
scheduler_discriminator = CosineAnnealingLR(optimizer_discriminator, T_max=epochs, eta_min=0.00001)

# 在训练循环中更新学习率
for epoch in range(epochs):
    # 训练代码...
    
    # 更新学习率
    if isinstance(scheduler_encoder, ReduceLROnPlateau):
        scheduler_encoder.step(val_rec_loss)
        scheduler_decoder.step(val_rec_loss)
    else:
        scheduler_encoder.step()
        scheduler_decoder.step()
        scheduler_discriminator.step()
    
    # 打印当前学习率
    print(f'Current learning rate: {optimizer_encoder.param_groups[0]["lr"]}')

# 在模型中使用批量归一化
class EncoderWithBatchNorm(nn.Module):
    def __init__(self, in_channels=3, out_channels=3):
        super(EncoderWithBatchNorm, self).__init__()
        
        # 下采样网络
        self.encoder = nn.Sequential(
            # 卷积层 + 批量归一化 + ReLU
            nn.Conv2d(in_channels * 2, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
        )
        
        # 上采样网络
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(64, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.Tanh()  # 输出范围[-1, 1]
        )
    
    def forward(self, cover, secret):
        # 连接载体图像和秘密图像
        x = torch.cat([cover, secret], dim=1)
        # 编码和解码
        x = self.encoder(x)
        x = self.decoder(x)
        # 将隐写图像与载体图像相加
        stego = cover + x
        return stego

# 在模型中使用层归一化
class EncoderWithLayerNorm(nn.Module):
    def __init__(self, in_channels=3, out_channels=3):
        super(EncoderWithLayerNorm, self).__init__()
        
        # 下采样网络
        self.encoder = nn.Sequential(
            # 卷积层 + 层归一化 + ReLU
            nn.Conv2d(in_channels * 2, 64, kernel_size=3, stride=2, padding=1),
            nn.LayerNorm([64, 128, 128]),  # 注意尺寸需要根据输入调整
            nn.ReLU(True),
            
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.LayerNorm([128, 64, 64]),
            nn.ReLU(True),
            
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.LayerNorm([256, 32, 32]),
            nn.ReLU(True),
        )
        
        # 上采样网络结构类似，省略...

# 权重初始化函数
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        # 卷积层使用正态分布初始化
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        if hasattr(m, 'bias') and m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
    elif classname.find('BatchNorm') != -1:
        # 批量归一化层初始化
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)
    elif classname.find('Linear') != -1:
        # 全连接层初始化
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

# 应用权重初始化
encoder.apply(weights_init)
decoder.apply(weights_init)
discriminator.apply(weights_init)

# 在训练循环中应用梯度裁剪
# 训练编码器和解码器
stego_images = encoder(cover_images, secret_images)
decoded_secret = decoder(stego_images)

reconstruction_loss = criterion_mse(decoded_secret, secret_images)
imperceptibility_loss = criterion_mse(stego_images, cover_images)
fake_outputs = discriminator(stego_images)
adversarial_loss = criterion_bce(fake_outputs, real_labels)

e_loss = alpha * reconstruction_loss + beta * imperceptibility_loss + gamma * adversarial_loss

# 反向传播
optimizer_encoder.zero_grad()
optimizer_decoder.zero_grad()
e_loss.backward()

# 应用梯度裁剪
torch.nn.utils.clip_grad_norm_(encoder.parameters(), max_norm=1.0)
torch.nn.utils.clip_grad_norm_(decoder.parameters(), max_norm=1.0)

# 参数更新
optimizer_encoder.step()
optimizer_decoder.step()

# 早停策略实现
patience = 20  # 容忍多少个epoch没有改善
best_val_loss = float('inf')
counter = 0

for epoch in range(epochs):
    # 训练和验证代码...
    
    # 检查是否有改善
    current_val_loss = val_rec_loss + val_imp_loss
    if current_val_loss < best_val_loss:
        best_val_loss = current_val_loss
        counter = 0
        # 保存最佳模型
        torch.save(encoder.state_dict(), 'best_encoder.pth')
        torch.save(decoder.state_dict(), 'best_decoder.pth')
        print(f'Saved best model at epoch {epoch+1}')
    else:
        counter += 1
        print(f'Early stopping counter: {counter}/{patience}')
        
        # 检查是否触发早停
        if counter >= patience:
            print(f'Early stopping at epoch {epoch+1}')
            break

# 超参数搜索空间
hyperparameters = {
    'batch_size': [16, 32, 64],
    'learning_rate': [0.001, 0.0005, 0.0002],
    'alpha': [1.0, 2.0, 5.0],  # 重建损失权重
    'beta': [1.0, 0.5, 0.1],   # 不可感知性损失权重
    'gamma': [0.1, 0.01, 0.001]  # 对抗损失权重
}

# 随机搜索示例
import random
import itertools

def random_search(hyperparameters, num_trials=20):
    # 生成所有可能的组合
    all_combinations = list(itertools.product(*hyperparameters.values()))
    # 随机选择指定数量的组合
    random.seed(42)
    selected_combinations = random.sample(all_combinations, min(num_trials, len(all_combinations)))
    
    best_score = float('inf')
    best_params = None
    
    for i, params in enumerate(selected_combinations):
        print(f'Trial {i+1}/{len(selected_combinations)}: {dict(zip(hyperparameters.keys(), params))}')
        
        # 设置超参数
        batch_size, lr, alpha, beta, gamma = params
        
        # 创建数据加载器
        train_loader, val_loader = get_data_loaders(train_dataset, val_dataset, batch_size=batch_size)
        
        # 初始化模型和优化器
        encoder = Encoder().to(device)
        decoder = Decoder().to(device)
        discriminator = Discriminator().to(device)
        
        optimizer_encoder = optim.Adam(encoder.parameters(), lr=lr)
        optimizer_decoder = optim.Adam(decoder.parameters(), lr=lr)
        optimizer_discriminator = optim.Adam(discriminator.parameters(), lr=lr)
        
        # 训练模型（使用较少的epoch进行验证）
        trial_best_loss = float('inf')
        for epoch in range(10):  # 只训练少量epoch进行验证
            # 训练代码...
            # 验证代码...
            
            current_val_loss = val_rec_loss + val_imp_loss
            if current_val_loss < trial_best_loss:
                trial_best_loss = current_val_loss
        
        # 更新最佳参数
        if trial_best_loss < best_score:
            best_score = trial_best_loss
            best_params = dict(zip(hyperparameters.keys(), params))
            print(f'New best score: {best_score:.4f} with params: {best_params}')
    
    return best_params, best_score

# 执行随机搜索
best_params, best_score = random_search(hyperparameters)
print(f'Best hyperparameters: {best_params}')
print(f'Best score: {best_score:.4f}')

# 损失函数权重调整示例
def find_optimal_weights(train_loader, val_loader, device):
    # 尝试不同的权重组合
    weights_to_try = [
        (1.0, 1.0, 0.1),  # alpha, beta, gamma
        (5.0, 1.0, 0.1),
        (1.0, 5.0, 0.1),
        (1.0, 1.0, 1.0),
        (5.0, 5.0, 0.1),
        (5.0, 1.0, 1.0),
        (1.0, 5.0, 1.0),
    ]
    
    best_weights = None
    best_val_loss = float('inf')
    
    for alpha, beta, gamma in weights_to_try:
        print(f'Trying weights: alpha={alpha}, beta={beta}, gamma={gamma}')
        
        # 初始化模型
        encoder = Encoder().to(device)
        decoder = Decoder().to(device)
        discriminator = Discriminator().to(device)
        
        # 初始化优化器
        optimizer_encoder = optim.Adam(encoder.parameters(), lr=0.0002)
        optimizer_decoder = optim.Adam(decoder.parameters(), lr=0.0002)
        optimizer_discriminator = optim.Adam(discriminator.parameters(), lr=0.0002)
        
        # 训练模型
        trial_best_loss = float('inf')
        for epoch in range(20):  # 训练少量epoch进行验证
            # 训练一个epoch
            train_e_loss, train_d_loss, train_disc_loss = train_one_epoch(
                encoder, decoder, discriminator, train_loader,
                optimizer_encoder, optimizer_decoder, optimizer_discriminator,
                criterion_mse, criterion_bce, device,
                alpha=alpha, beta=beta, gamma=gamma
            )
            
            # 验证
            val_rec_loss, val_imp_loss = validate(
                encoder, decoder, val_loader, criterion_mse, device
            )
            
            current_val_loss = val_rec_loss + val_imp_loss
            if current_val_loss < trial_best_loss:
                trial_best_loss = current_val_loss
        
        # 更新最佳权重
        if trial_best_loss < best_val_loss:
            best_val_loss = trial_best_loss
            best_weights = (alpha, beta, gamma)
            print(f'New best weights: {best_weights} with loss: {best_val_loss:.4f}')
    
    return best_weights, best_val_loss

# 寻找最佳损失函数权重
best_weights, best_val_loss = find_optimal_weights(train_loader, val_loader, device)
alpha, beta, gamma = best_weights
print(f'Best loss weights: alpha={alpha}, beta={beta}, gamma={gamma}')

# 可视化训练和验证损失
def plot_loss_history(loss_history):
    plt.figure(figsize=(15, 10))
    
    # 绘制训练损失
    plt.subplot(2, 2, 1)
    plt.plot(loss_history['encoder'], label='Encoder Loss')
    plt.plot(loss_history['decoder'], label='Decoder Loss')
    plt.plot(loss_history['discriminator'], label='Discriminator Loss')
    plt.title('Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    
    # 绘制验证损失
    plt.subplot(2, 2, 2)
    plt.plot(loss_history['val_reconstruction'], label='Reconstruction Loss')
    plt.plot(loss_history['val_imperceptibility'], label='Imperceptibility Loss')
    plt.title('Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    
    # 绘制总验证损失
    total_val_loss = [loss_history['val_reconstruction'][i] + loss_history['val_imperceptibility'][i] 
                      for i in range(len(loss_history['val_reconstruction']))]
    plt.subplot(2, 2, 3)
    plt.plot(total_val_loss, label='Total Validation Loss')
    plt.title('Total Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    
    # 绘制学习率变化（如果记录了）
    if 'learning_rate' in loss_history:
        plt.subplot(2, 2, 4)
        plt.plot(loss_history['learning_rate'], label='Learning Rate')
        plt.title('Learning Rate Schedule')
        plt.xlabel('Epoch')
        plt.ylabel('Learning Rate')
        plt.yscale('log')
        plt.legend()
        plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('loss_history.png')
    plt.show()

# 调用函数绘制损失曲线
plot_loss_history(loss_history)

# 可视化隐写结果
def visualize_results(encoder, decoder, val_loader, device, num_samples=5):
    encoder.eval()
    decoder.eval()
    
    # 获取验证集中的样本
    with torch.no_grad():
        # 获取一个批次的数据
        for i, data in enumerate(val_loader):
            cover_images = data['cover'].to(device)
            secret_images = data['secret'].to(device)
            
            # 生成隐写图像和解码秘密图像
            stego_images = encoder(cover_images, secret_images)
            decoded_secret = decoder(stego_images)
            
            # 转换为numpy数组并反归一化
            cover_images = cover_images.cpu().numpy()
            secret_images = secret_images.cpu().numpy()
            stego_images = stego_images.cpu().numpy()
            decoded_secret = decoded_secret.cpu().numpy()
            
            # 反归一化（从[-1, 1]到[0, 1]）
            cover_images = (cover_images + 1) / 2
            secret_images = (secret_images + 1) / 2
            stego_images = (stego_images + 1) / 2
            decoded_secret = (decoded_secret + 1) / 2
            
            # 绘制结果
            plt.figure(figsize=(20, num_samples * 5))
            
            for j in range(num_samples):
                # 原始载体图像
                plt.subplot(num_samples, 4, j*4 + 1)
                plt.imshow(np.transpose(cover_images[j], (1, 2, 0)))
                plt.title(f'Cover Image {j+1}')
                plt.axis('off')
                
                # 秘密图像
                plt.subplot(num_samples, 4, j*4 + 2)
                plt.imshow(np.transpose(secret_images[j], (1, 2, 0)))
                plt.title(f'Secret Image {j+1}')
                plt.axis('off')
                
                # 隐写图像
                plt.subplot(num_samples, 4, j*4 + 3)
                plt.imshow(np.transpose(stego_images[j], (1, 2, 0)))
                plt.title(f'Stego Image {j+1}')
                plt.axis('off')
                
                # 解码秘密图像
                plt.subplot(num_samples, 4, j*4 + 4)
                plt.imshow(np.transpose(decoded_secret[j], (1, 2, 0)))
                plt.title(f'Decoded Secret {j+1}')
                plt.axis('off')
            
            plt.tight_layout()
            plt.savefig('steganography_results.png')
            plt.show()
            
            # 只处理第一个批次
            break

# 可视化隐写结果
visualize_results(encoder, decoder, val_loader, device)

# 计算模型评估指标
def evaluate_model(encoder, decoder, test_loader, device):
    encoder.eval()
    decoder.eval()
    
    total_mse_reconstruction = 0
    total_mse_imperceptibility = 0
    total_psnr_cover_stego = 0
    total_psnr_secret_decoded = 0
    total_ssim_cover_stego = 0
    total_ssim_secret_decoded = 0
    
    from skimage.metrics import structural_similarity as ssim
    
    with torch.no_grad():
        for data in test_loader:
            cover_images = data['cover'].to(device)
            secret_images = data['secret'].to(device)
            
            # 生成隐写图像和解码秘密图像
            stego_images = encoder(cover_images, secret_images)
            decoded_secret = decoder(stego_images)
            
            # 转换为numpy数组并反归一化
            cover_images_np = cover_images.cpu().numpy()
            secret_images_np = secret_images.cpu().numpy()
            stego_images_np = stego_images.cpu().numpy()
            decoded_secret_np = decoded_secret.cpu().numpy()
            
            # 反归一化（从[-1, 1]到[0, 1]）
            cover_images_np = (cover_images_np + 1) / 2
            secret_images_np = (secret_images_np + 1) / 2
            stego_images_np = (stego_images_np + 1) / 2
            decoded_secret_np = (decoded_secret_np + 1) / 2
            
            # 计算每对图像的指标
            for i in range(cover_images.size(0)):
                # 计算MSE
                mse_reconstruction = np.mean((secret_images_np[i] - decoded_secret_np[i]) ** 2)
                mse_imperceptibility = np.mean((cover_images_np[i] - stego_images_np[i]) ** 2)
                
                # 计算PSNR
                max_pixel = 1.0
                psnr_cover_stego = 20 * np.log10(max_pixel / np.sqrt(mse_imperceptibility))
                psnr_secret_decoded = 20 * np.log10(max_pixel / np.sqrt(mse_reconstruction))
                
                # 计算SSIM（需要转换为2D图像）
                cover_2d = np.mean(cover_images_np[i].transpose(1, 2, 0), axis=2)
                stego_2d = np.mean(stego_images_np[i].transpose(1, 2, 0), axis=2)
                secret_2d = np.mean(secret_images_np[i].transpose(1, 2, 0), axis=2)
                decoded_2d = np.mean(decoded_secret_np[i].transpose(1, 2, 0), axis=2)
                
                ssim_cover_stego = ssim(cover_2d, stego_2d, data_range=1.0)
                ssim_secret_decoded = ssim(secret_2d, decoded_2d, data_range=1.0)
                
                # 累加指标
                total_mse_reconstruction += mse_reconstruction
                total_mse_imperceptibility += mse_imperceptibility
                total_psnr_cover_stego += psnr_cover_stego
                total_psnr_secret_decoded += psnr_secret_decoded
                total_ssim_cover_stego += ssim_cover_stego
                total_ssim_secret_decoded += ssim_secret_decoded
    
    # 计算平均指标
    num_samples = len(test_loader.dataset)
    metrics = {
        'mse_reconstruction': total_mse_reconstruction / num_samples,
        'mse_imperceptibility': total_mse_imperceptibility / num_samples,
        'psnr_cover_stego': total_psnr_cover_stego / num_samples,
        'psnr_secret_decoded': total_psnr_secret_decoded / num_samples,
        'ssim_cover_stego': total_ssim_cover_stego / num_samples,
        'ssim_secret_decoded': total_ssim_secret_decoded / num_samples
    }
    
    return metrics

# 评估模型性能
test_metrics = evaluate_model(encoder, decoder, test_loader, device)
print("Test Metrics:")
for key, value in test_metrics.items():
    print(f"{key}: {value:.4f}")

# 使用残差连接的编码器示例
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        # 快捷连接
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += self.shortcut(x)  # 残差连接
        out = self.relu(out)
        return out

# 使用残差块的编码器
class EncoderWithResidual(nn.Module):
    def __init__(self, in_channels=3, out_channels=3):
        super(EncoderWithResidual, self).__init__()
        # 初始卷积
        self.initial_conv = nn.Sequential(
            nn.Conv2d(in_channels * 2, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        
        # 残差块
        self.layer1 = ResidualBlock(64, 64)
        self.layer2 = ResidualBlock(64, 128, stride=2)
        self.layer3 = ResidualBlock(128, 256, stride=2)
        
        # 上采样
        self.upsample = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            
            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(64, out_channels, kernel_size=3, stride=1, padding=1),
            nn.Tanh()
        )
    
    def forward(self, cover, secret):
        x = torch.cat([cover, secret], dim=1)
        x = self.initial_conv(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.upsample(x)
        stego = cover + x
        return stego

# 小批量判别示例
class MinibatchDiscrimination(nn.Module):
    def __init__(self, in_features, out_features, kernel_dims, mean=True):
        super(MinibatchDiscrimination, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.kernel_dims = kernel_dims
        self.mean = mean
        self.T = nn.Parameter(torch.Tensor(in_features, out_features, kernel_dims))
        nn.init.normal_(self.T, 0, 1)
    
    def forward(self, x):
        # x: [N, in_features]
        M = torch.mm(x, self.T.view(self.in_features, -1))
        # M: [N, out_features * kernel_dims]
        M = M.view(-1, self.out_features, self.kernel_dims)
        # M: [N, out_features, kernel_dims]
        
        # 计算实例之间的L1距离
        op1 = M.unsqueeze(0)  # [1, N, out_features, kernel_dims]
        op2 = M.unsqueeze(1)  # [N, 1, out_features, kernel_dims]
        # 计算L1距离
        abs_diff = torch.sum(torch.abs(op1 - op2), dim=3)
        # 计算相似度
        features = torch.sum(torch.exp(-abs_diff), dim=1)
        # features: [N, out_features]
        
        # 添加到原始特征
        if self.mean:
            features = features / x.size(0)
        
        return torch.cat([x, features], 1)

# 在判别器中使用小批量判别
class DiscriminatorWithMinibatch(nn.Module):
    def __init__(self, in_channels=3):
        super(DiscriminatorWithMinibatch, self).__init__()
        
        self.features = nn.Sequential(
            # 卷积层
            nn.Conv2d(in_channels, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
        )
        
        # 小批量判别
        self.mbd = MinibatchDiscrimination(512 * 16 * 16, 64, 8)
        
        # 输出层
        self.output = nn.Sequential(
            nn.Linear(512 * 16 * 16 + 64, 1),
            # 不使用sigmoid，因为使用了BCEWithLogitsLoss
        )
    
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.mbd(x)
        x = self.output(x)
        return x

# 在模型中添加Dropout
class EncoderWithDropout(nn.Module):
    def __init__(self, in_channels=3, out_channels=3, dropout_rate=0.5):
        super(EncoderWithDropout, self).__init__()
        
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels * 2, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.Dropout(dropout_rate),  # 添加Dropout
            
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.Dropout(dropout_rate),  # 添加Dropout
            
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
        )
        
        # 上采样网络结构类似，省略...

# 使用L2正则化的优化器
optimizer_encoder = optim.Adam(encoder.parameters(), lr=0.0002, weight_decay=1e-5)  # 添加L2正则化
optimizer_decoder = optim.Adam(decoder.parameters(), lr=0.0002, weight_decay=1e-5)

# 使用AdamW优化器
optimizer_encoder = optim.AdamW(encoder.parameters(), lr=0.0002, weight_decay=1e-4)
optimizer_decoder = optim.AdamW(decoder.parameters(), lr=0.0002, weight_decay=1e-4)
optimizer_discriminator = optim.AdamW(discriminator.parameters(), lr=0.0002, weight_decay=1e-4)

# 梯度累积示例
accumulation_steps = 4  # 累积4个批次的梯度

for i, data in enumerate(train_loader):
    cover_images = data['cover'].to(device)
    secret_images = data['secret'].to(device)
    
    # 前向传播
    stego_images = encoder(cover_images, secret_images)
    decoded_secret = decoder(stego_images)
    
    # 计算损失
    reconstruction_loss = criterion_mse(decoded_secret, secret_images)
    imperceptibility_loss = criterion_mse(stego_images, cover_images)
    
    # 缩放损失以补偿累积
    loss = (reconstruction_loss + imperceptibility_loss) / accumulation_steps
    
    # 反向传播
    loss.backward()
    
    # 只有在累积了足够的批次后才更新参数
    if (i + 1) % accumulation_steps == 0:
        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(encoder.parameters(), max_norm=1.0)
        torch.nn.utils.clip_grad_norm_(decoder.parameters(), max_norm=1.0)
        
        # 参数更新
        optimizer_encoder.step()
        optimizer_decoder.step()
        
        # 梯度清零
        optimizer_encoder.zero_grad()
        optimizer_decoder.zero_grad()

# 加载训练好的模型
encoder = Encoder().to(device)
decoder = Decoder().to(device)

encoder.load_state_dict(torch.load('best_encoder.pth', map_location=device))
decoder.load_state_dict(torch.load('best_decoder.pth', map_location=device))

# 设置为评估模式
encoder.eval()
decoder.eval()

# 导出为ONNX格式
def export_to_onnx(model, dummy_input, output_path, dynamic_axes=None):
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        export_params=True,
        opset_version=11,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes=dynamic_axes if dynamic_axes else {}
    )
    print(f'Model exported to {output_path}')

# 准备虚拟输入
dummy_cover = torch.randn(1, 3, 256, 256).to(device)  # 假设输入图像大小为256x256
dummy_secret = torch.randn(1, 3, 256, 256).to(device)

# 导出编码器（需要修改以接受两个输入）
class EncoderWrapper(nn.Module):
    def __init__(self, encoder):
        super(EncoderWrapper, self).__init__()
        self.encoder = encoder
    
    def forward(self, x):
        # 假设x的前3个通道是载体图像，后3个通道是秘密图像
        cover = x[:, :3, :, :]
        secret = x[:, 3:, :, :]
        return self.encoder(cover, secret)

# 包装编码器
en_wrapper = EncoderWrapper(encoder).to(device)
dummy_input_combined = torch.cat([dummy_cover, dummy_secret], dim=1)

# 导出为ONNX
export_to_onnx(
    en_wrapper,
    dummy_input_combined,
    'encoder.onnx',
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size', 2: 'height', 3: 'width'}
    }
)

# 导出解码器
export_to_onnx(
    decoder,
    dummy_cover,
    'decoder.onnx',
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size', 2: 'height', 3: 'width'}
    }
)

# 动态量化
def quantize_dynamic(model, example_input, output_path):
    # 将模型转换为CPU模式
    model_cpu = model.cpu()
    # 动态量化
    quantized_model = torch.quantization.quantize_dynamic(
        model_cpu,
        {nn.Linear, nn.Conv2d, nn.ConvTranspose2d},
        dtype=torch.qint8
    )
    # 保存量化后的模型
    torch.jit.save(torch.jit.script(quantized_model), output_path)
    print(f'Quantized model saved to {output_path}')

# 量化编码器和解码器
encoder_cpu = encoder.cpu()
decoder_cpu = decoder.cpu()

example_cover = torch.randn(1, 3, 256, 256)
example_secret = torch.randn(1, 3, 256, 256)

# 对于编码器，我们需要一个包装器来处理两个输入
class QuantizableEncoderWrapper(nn.Module):
    def __init__(self, encoder):
        super(QuantizableEncoderWrapper, self).__init__()
        self.encoder = encoder
    
    def forward(self, cover, secret):
        return self.encoder(cover, secret)

# 包装并量化编码器
en_wrapper_quant = QuantizableEncoderWrapper(encoder_cpu)
torch.jit.script(en_wrapper_quant, (example_cover, example_secret))
torch.save(en_wrapper_quant.state_dict(), 'quantizable_encoder_wrapper.pth')

# 量化解码器
torch.jit.script(decoder_cpu, example_cover)
torch.save(decoder_cpu.state_dict(), 'quantizable_decoder.pth')

# 动态量化示例（注意：完整的量化流程可能需要更多配置）
# 这里仅提供基本示例
quantize_dynamic(decoder_cpu, example_cover, 'decoder_quantized.pt')

# 模型剪枝示例
import torch.nn.utils.prune as prune

# 加载模型
encoder_prunable = Encoder().to(device)
encoder_prunable.load_state_dict(torch.load('best_encoder.pth', map_location=device))

# 对卷积层应用L1范数剪枝
parameters_to_prune = []
for name, module in encoder_prunable.named_modules():
    if isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d):
        parameters_to_prune.append((module, 'weight'))

# 全局剪枝：移除10%的权重
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.1,
)

# 使剪枝永久化
for module, name in parameters_to_prune:
    prune.remove(module, name)

# 保存剪枝后的模型
torch.save(encoder_prunable.state_dict(), 'encoder_pruned.pth')

# 验证剪枝后的模型性能
# 这里需要与原始模型进行比较测试

# 使用TorchScript优化模型
def optimize_with_torchscript(model, example_inputs, output_path):
    # 跟踪模型
    traced_model = torch.jit.trace(model, example_inputs)
    # 优化模型
    optimized_model = torch.jit.optimize_for_inference(traced_model)
    # 保存优化后的模型
    torch.jit.save(optimized_model, output_path)
    print(f'Optimized model saved to {output_path}')
    return optimized_model

# 优化编码器
class EncoderScriptWrapper(nn.Module):
    def __init__(self, encoder):
        super(EncoderScriptWrapper, self).__init__()
        self.encoder = encoder
    
    def forward(self, cover, secret):
        return self.encoder(cover, secret)

# 准备示例输入
example_cover = torch.randn(1, 3, 256, 256).to(device)
example_secret = torch.randn(1, 3, 256, 256).to(device)

# 包装并优化编码器
en_script_wrapper = EncoderScriptWrapper(encoder).to(device)
traced_encoder = optimize_with_torchscript(
    en_script_wrapper,
    (example_cover, example_secret),
    'encoder_traced.pt'
)

# 优化解码器
traced_decoder = optimize_with_torchscript(
    decoder,
    example_cover,
    'decoder_traced.pt'
)

# 使用Flask部署模型
from flask import Flask, request, jsonify
import base64
from PIL import Image
import io
import numpy as np

def preprocess_image(image_bytes, target_size=(256, 256)):
    # 将图像字节转换为PIL图像
    image = Image.open(io.BytesIO(image_bytes))
    # 调整大小
    image = image.resize(target_size)
    # 转换为numpy数组
    image = np.array(image)
    # 归一化到[-1, 1]
    image = (image / 127.5) - 1.0
    # 添加批次维度
    image = np.transpose(image, (2, 0, 1))
    image = np.expand_dims(image, axis=0)
    # 转换为tensor
    return torch.from_numpy(image).float().to(device)

def postprocess_image(tensor):
    # 将tensor转换为numpy数组
    image = tensor.squeeze().cpu().detach().numpy()
    # 反归一化到[0, 1]
    image = (image + 1.0) / 2.0
    # 转换为[0, 255]范围
    image = (image * 255).astype(np.uint8)
    # 调整通道顺序
    image = np.transpose(image, (1, 2, 0))
    # 转换为PIL图像
    return Image.fromarray(image)

# 初始化Flask应用
app = Flask(__name__)

@app.route('/encode', methods=['POST'])
def encode():
    try:
        # 获取请求数据
        data = request.json
        cover_base64 = data['cover_image']
        secret_base64 = data['secret_image']
        
        # 解码base64图像数据
        cover_bytes = base64.b64decode(cover_base64)
        secret_bytes = base64.b64decode(secret_base64)
        
        # 预处理图像
        cover_tensor = preprocess_image(cover_bytes)
        secret_tensor = preprocess_image(secret_bytes)
        
        # 执行隐写编码
        with torch.no_grad():
            stego_tensor = encoder(cover_tensor, secret_tensor)
        
        # 后处理隐写图像
        stego_image = postprocess_image(stego_tensor)
        
        # 将图像转换为base64
        buffer = io.BytesIO()
        stego_image.save(buffer, format='PNG')
        buffer.seek(0)
        stego_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        # 返回结果
        return jsonify({
            'status': 'success',
            'stego_image': stego_base64
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 500

@app.route('/decode', methods=['POST'])
def decode():
    try:
        # 获取请求数据
        data = request.json
        stego_base64 = data['stego_image']
        
        # 解码base64图像数据
        stego_bytes = base64.b64decode(stego_base64)
        
        # 预处理图像
        stego_tensor = preprocess_image(stego_bytes)
        
        # 执行隐写解码
        with torch.no_grad():
            secret_tensor = decoder(stego_tensor)
        
        # 后处理秘密图像
        secret_image = postprocess_image(secret_tensor)
        
        # 将图像转换为base64
        buffer = io.BytesIO()
        secret_image.save(buffer, format='PNG')
        buffer.seek(0)
        secret_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        # 返回结果
        return jsonify({
            'status': 'success',
            'secret_image': secret_base64
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 500

if __name__ == '__main__':
    # 加载模型
    encoder = Encoder().to(device)
    decoder = Decoder().to(device)
    encoder.load_state_dict(torch.load('best_encoder.pth', map_location=device))
    decoder.load_state_dict(torch.load('best_decoder.pth', map_location=device))
    encoder.eval()
    decoder.eval()
    
    # 启动服务器
    app.run(host='0.0.0.0', port=5000)

# 使用PyInstaller打包为可执行文件的示例代码
# 这个脚本将作为应用程序的入口点

import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import torch
import numpy as np
import io
import os

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义模型类
class Encoder(torch.nn.Module):
    # 与训练时相同的模型定义
    pass

class Decoder(torch.nn.Module):
    # 与训练时相同的模型定义
    pass

# 加载模型
def load_models():
    encoder = Encoder().to(device)
    decoder = Decoder().to(device)
    
    # 获取当前脚本所在目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 加载模型权重
    encoder_path = os.path.join(script_dir, 'models', 'encoder_traced.pt')
    decoder_path = os.path.join(script_dir, 'models', 'decoder_traced.pt')
    
    encoder = torch.jit.load(encoder_path, map_location=device)
    decoder = torch.jit.load(decoder_path, map_location=device)
    
    encoder.eval()
    decoder.eval()
    
    return encoder, decoder

# 图像处理函数
def preprocess_image(image_path, target_size=(256, 256)):
    image = Image.open(image_path).convert('RGB')
    image = image.resize(target_size)
    image_np = np.array(image)
    image_np = (image_np / 127.5) - 1.0  # 归一化到[-1, 1]
    image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).float().unsqueeze(0).to(device)
    return image_tensor, image_np.shape

def postprocess_image(tensor, original_shape=None):
    image = tensor.squeeze().cpu().detach().numpy()
    image = (image + 1.0) / 2.0  # 反归一化到[0, 1]
    image = (image * 255).astype(np.uint8)
    image = np.transpose(image, (1, 2, 0))
    
    # 如果提供了原始形状，则调整大小
    if original_shape is not None:
        image = Image.fromarray(image).resize((original_shape[1], original_shape[0]))
        image = np.array(image)
    
    return Image.fromarray(image)

# GUI应用类
class SteganographyApp:
    def __init__(self, root):
        self.root = root
        self.root.title("机器学习隐写工具")
        self.root.geometry("800x600")
        
        # 加载模型
        try:
            self.encoder, self.decoder = load_models()
        except Exception as e:
            messagebox.showerror("错误", f"加载模型失败: {str(e)}")
            self.root.destroy()
            return
        
        # 变量初始化
        self.cover_path = ""
        self.secret_path = ""
        self.stego_path = ""
        
        # 创建UI
        self.create_ui()
    
    def create_ui(self):
        # 创建选项卡
        tab_control = tk.ttk.Notebook(self.root)
        
        # 编码选项卡
        encode_tab = tk.Frame(tab_control)
        tab_control.add(encode_tab, text="编码")
        
        # 解码选项卡
        decode_tab = tk.Frame(tab_control)
        tab_control.add(decode_tab, text="解码")
        
        tab_control.pack(expand=1, fill="both")
        
        # 设置编码选项卡
        self.setup_encode_tab(encode_tab)
        
        # 设置解码选项卡
        self.setup_decode_tab(decode_tab)
    
    def setup_encode_tab(self, parent):
        # 载体图像选择
        tk.Label(parent, text="选择载体图像:").grid(row=0, column=0, padx=10, pady=10)
        self.cover_label = tk.Label(parent, text="未选择")
        self.cover_label.grid(row=0, column=1, padx=10, pady=10)
        tk.Button(parent, text="浏览", command=self.select_cover).grid(row=0, column=2, padx=10, pady=10)
        
        # 秘密图像选择
        tk.Label(parent, text="选择秘密图像:").grid(row=1, column=0, padx=10, pady=10)
        self.secret_label = tk.Label(parent, text="未选择")
        self.secret_label.grid(row=1, column=1, padx=10, pady=10)
        tk.Button(parent, text="浏览", command=self.select_secret).grid(row=1, column=2, padx=10, pady=10)
        
        # 编码按钮
        tk.Button(parent, text="开始编码", command=self.encode).grid(row=2, column=0, columnspan=3, pady=20)
        
        # 结果保存位置
        tk.Label(parent, text="隐写图像保存为:").grid(row=3, column=0, padx=10, pady=10)
        self.stego_label = tk.Label(parent, text="未保存")
        self.stego_label.grid(row=3, column=1, padx=10, pady=10)
    
    def setup_decode_tab(self, parent):
        # 隐写图像选择
        tk.Label(parent, text="选择隐写图像:").grid(row=0, column=0, padx=10, pady=10)
        self.decode_stego_label = tk.Label(parent, text="未选择")
        self.decode_stego_label.grid(row=0, column=1, padx=10, pady=10)
        tk.Button(parent, text="浏览", command=self.select_stego_for_decode).grid(row=0, column=2, padx=10, pady=10)
        
        # 解码按钮
        tk.Button(parent, text="开始解码", command=self.decode).grid(row=1, column=0, columnspan=3, pady=20)
        
        # 结果保存位置
        tk.Label(parent, text="解码图像保存为:").grid(row=2, column=0, padx=10, pady=10)
        self.decoded_label = tk.Label(parent, text="未保存")
        self.decoded_label.grid(row=2, column=1, padx=10, pady=10)
    
    def select_cover(self):
        path = filedialog.askopenfilename(filetypes=[("Image files", "*.png;*.jpg;*.jpeg")])
        if path:
            self.cover_path = path
            self.cover_label.config(text=os.path.basename(path))
    
    def select_secret(self):
        path = filedialog.askopenfilename(filetypes=[("Image files", "*.png;*.jpg;*.jpeg")])
        if path:
            self.secret_path = path
            self.secret_label.config(text=os.path.basename(path))
    
    def select_stego_for_decode(self):
        path = filedialog.askopenfilename(filetypes=[("Image files", "*.png;*.jpg;*.jpeg")])
        if path:
            self.stego_path = path
            self.decode_stego_label.config(text=os.path.basename(path))
    
    def encode(self):
        if not self.cover_path or not self.secret_path:
            messagebox.showerror("错误", "请选择载体图像和秘密图像")
            return
        
        try:
            # 预处理图像
            cover_tensor, cover_shape = preprocess_image(self.cover_path)
            secret_tensor, _ = preprocess_image(self.secret_path)
            
            # 执行编码
            with torch.no_grad():
                stego_tensor = self.encoder(cover_tensor, secret_tensor)
            
            # 后处理并保存图像
            stego_image = postprocess_image(stego_tensor, cover_shape)
            
            # 保存文件
            save_path = filedialog.asksaveasfilename(
                defaultextension=".png",
                filetypes=[("PNG files", "*.png")]
            )
            
            if save_path:
                stego_image.save(save_path)
                self.stego_label.config(text=os.path.basename(save_path))
                messagebox.showinfo("成功", "编码完成！")
        except Exception as e:
            messagebox.showerror("错误", f"编码失败: {str(e)}")
    
    def decode(self):
        if not self.stego_path:
            messagebox.showerror("错误", "请选择隐写图像")
            return
        
        try:
            # 预处理图像
            stego_tensor, stego_shape = preprocess_image(self.stego_path)
            
            # 执行解码
            with torch.no_grad():
                secret_tensor = self.decoder(stego_tensor)
            
            # 后处理并保存图像
            secret_image = postprocess_image(secret_tensor, stego_shape)
            
            # 保存文件
            save_path = filedialog.asksaveasfilename(
                defaultextension=".png",
                filetypes=[("PNG files", "*.png")]
            )
            
            if save_path:
                secret_image.save(save_path)
                self.decoded_label.config(text=os.path.basename(save_path))
                messagebox.showinfo("成功", "解码完成！")
        except Exception as e:
            messagebox.showerror("错误", f"解码失败: {str(e)}")

# 主函数
if __name__ == "__main__":
    root = tk.Tk()
    app = SteganographyApp(root)
    root.mainloop()

# PyTorch Mobile 部署示例
# 1. 首先将模型转换为移动友好格式
def prepare_for_mobile(model, example_input, output_path):
    # 跟踪模型
    traced_model = torch.jit.trace(model, example_input)
    # 针对移动设备优化
    mobile_model = traced_model._save_for_lite_interpreter(output_path)
    print(f'Mobile model saved to {output_path}')
    return mobile_model

# 准备移动设备模型
example_cover = torch.randn(1, 3, 256, 256).to(device)
example_secret = torch.randn(1, 3, 256, 256).to(device)

# 为移动设备准备编码器（使用包装器）
en_mobile_wrapper = EncoderScriptWrapper(encoder).to(device)
prepare_for_mobile(
    en_mobile_wrapper,
    (example_cover, example_secret),
    'encoder_mobile.ptl'
)

# 为移动设备准备解码器
prepare_for_mobile(
    decoder,
    example_cover,
    'decoder_mobile.ptl'
)

# 数字水印嵌入示例
def embed_watermark(image_path, watermark_text, output_path):
    # 加载图像
    cover_image = Image.open(image_path).convert('RGB')
    
    # 将文本水印转换为图像
    from PIL import ImageDraw, ImageFont
    watermark_size = (128, 64)
    watermark_img = Image.new('RGB', watermark_size, color='white')
    d = ImageDraw.Draw(watermark_img)
    try:
        # 使用默认字体
        font = ImageFont.load_default()
    except:
        # 如果没有默认字体，使用简单文本
        font = None
    
    # 在水印图像上绘制文本
    d.text((10, 10), watermark_text, fill='black', font=font)
    
    # 调整大小
    cover_resized = cover_image.resize((256, 256))
    watermark_resized = watermark_img.resize((256, 256))
    
    # 预处理
    cover_tensor = preprocess_pil_image(cover_resized)
    watermark_tensor = preprocess_pil_image(watermark_resized)
    
    # 嵌入水印
    with torch.no_grad():
        watermarked_tensor = encoder(cover_tensor, watermark_tensor)
    
    # 后处理
    watermarked_image = postprocess_tensor(watermarked_tensor, cover_image.size)
    
    # 保存
    watermarked_image.save(output_path)
    print(f'Watermarked image saved to {output_path}')
    
    return output_path

def extract_watermark(watermarked_path, output_path):
    # 加载图像
    watermarked_image = Image.open(watermarked_path).convert('RGB')
    
    # 预处理
    watermarked_tensor = preprocess_pil_image(watermarked_image)
    
    # 提取水印
    with torch.no_grad():
        extracted_tensor = decoder(watermarked_tensor)
    
    # 后处理
    extracted_image = postprocess_tensor(extracted_tensor)
    
    # 保存
    extracted_image.save(output_path)
    print(f'Extracted watermark saved to {output_path}')
    
    return output_path

# 文本隐写通信示例
def text_to_image(text, image_size=(256, 256)):
    """将文本转换为图像"""
    # 创建空白图像
    image = Image.new('RGB', image_size, color='white')
    draw = ImageDraw.Draw(image)
    
    # 将文本转换为二进制
    binary_text = ''.join(format(ord(char), '08b') for char in text)
    
    # 确保文本不会超过图像容量
    max_chars = (image_size[0] * image_size[1] * 3) // 8
    if len(binary_text) > max_chars * 8:
        raise ValueError(f"文本太长，最大支持 {max_chars} 个字符")
    
    # 将二进制数据填充到图像中
    pixels = image.load()
    index = 0
    
    for i in range(image_size[0]):
        for j in range(image_size[1]):
            r, g, b = pixels[i, j]
            
            # 修改最低有效位
            if index < len(binary_text):
                r = (r & ~1) | int(binary_text[index])
                index += 1
            if index < len(binary_text):
                g = (g & ~1) | int(binary_text[index])
                index += 1
            if index < len(binary_text):
                b = (b & ~1) | int(binary_text[index])
                index += 1
            
            pixels[i, j] = (r, g, b)
    
    return image

def secure_communication(cover_image_path, secret_text, output_path):
    # 将文本转换为图像
    secret_image = text_to_image(secret_text)
    
    # 保存秘密图像（临时）
    temp_secret_path = "temp_secret.png"
    secret_image.save(temp_secret_path)
    
    try:
        # 使用机器学习隐写技术嵌入秘密信息
        cover_image = Image.open(cover_image_path).convert('RGB')
        cover_resized = cover_image.resize((256, 256))
        
        # 预处理
        cover_tensor = preprocess_pil_image(cover_resized)
        secret_tensor = preprocess_pil_image(secret_image)
        
        # 嵌入
        with torch.no_grad():
            stego_tensor = encoder(cover_tensor, secret_tensor)
        
        # 后处理
        stego_image = postprocess_tensor(stego_tensor, cover_image.size)
        stego_image.save(output_path)
        print(f"安全通信图像已保存到 {output_path}")
        
    finally:
        # 删除临时文件
        if os.path.exists(temp_secret_path):
            os.remove(temp_secret_path)

def receive_secure_message(stego_image_path, output_text_path=None):
    # 加载隐写图像
    stego_image = Image.open(stego_image_path).convert('RGB')
    
    # 预处理
    stego_tensor = preprocess_pil_image(stego_image)
    
    # 提取秘密图像
    with torch.no_grad():
        secret_tensor = decoder(stego_tensor)
    
    # 后处理
    secret_image = postprocess_tensor(secret_tensor)
    
    # 从图像中提取文本
    pixels = secret_image.load()
    binary_text = ""
    
    # 读取最低有效位
    for i in range(secret_image.width):
        for j in range(secret_image.height):
            r, g, b = pixels[i, j]
            binary_text += str(r & 1)
            binary_text += str(g & 1)
            binary_text += str(b & 1)
    
    # 将二进制转换为文本
    text = ""
    for i in range(0, len(binary_text), 8):
        byte = binary_text[i:i+8]
        if len(byte) < 8:
            break
        char = chr(int(byte, 2))
        text += char
        
        # 如果遇到结束符，停止解析
        if char == '\0':
            break
    
    # 保存文本
    if output_text_path:
        with open(output_text_path, 'w', encoding='utf-8') as f:
            f.write(text)
        print(f"解密文本已保存到 {output_text_path}")
    
    return text

# 隐私保护示例
def privacy_protection(image_path, regions_to_hide, output_path):
    """
    使用隐写技术隐藏图像中的敏感区域
    
    参数:
    - image_path: 原始图像路径
    - regions_to_hide: 需要隐藏的区域列表，每个区域格式为 (x, y, width, height)
    - output_path: 输出图像路径
    """
    # 加载图像
    original_image = Image.open(image_path).convert('RGB')
    
    # 创建秘密图像，只包含需要隐藏的区域
    secret_image = Image.new('RGB', original_image.size, color='white')
    secret_draw = ImageDraw.Draw(secret_image)
    
    # 在原始图像上标记需要隐藏的区域
    modified_image = original_image.copy()
    draw = ImageDraw.Draw(modified_image)
    
    for region in regions_to_hide:
        x, y, w, h = region
        # 从原始图像复制区域到秘密图像
        region_img = original_image.crop((x, y, x + w, y + h))
        secret_image.paste(region_img, (x, y))
        
        # 在原始图像上用模糊或其他方式替换该区域
        # 这里简化处理，使用白色矩形
        draw.rectangle([x, y, x + w, y + h], fill='white')
    
    # 调整大小以适应模型输入
    modified_resized = modified_image.resize((256, 256))
    secret_resized = secret_image.resize((256, 256))
    
    # 预处理
    modified_tensor = preprocess_pil_image(modified_resized)
    secret_tensor = preprocess_pil_image(secret_resized)
    
    # 执行隐写
    with torch.no_grad():
        stego_tensor = encoder(modified_tensor, secret_tensor)
    
    # 后处理
    stego_image = postprocess_tensor(stego_tensor, original_image.size)
    stego_image.save(output_path)
    
    print(f"隐私保护后的图像已保存到 {output_path}")
    return output_path

def recover_privacy(stego_image_path, original_size, output_path):
    """恢复被隐藏的隐私信息"""
    # 加载隐写图像
    stego_image = Image.open(stego_image_path).convert('RGB')
    
    # 预处理
    stego_tensor = preprocess_pil_image(stego_image)
    
    # 提取秘密信息
    with torch.no_grad():
        secret_tensor = decoder(stego_tensor)
    
    # 后处理
    secret_image = postprocess_tensor(secret_tensor, original_size)
    secret_image.save(output_path)
    
    print(f"恢复的隐私信息已保存到 {output_path}")
    return output_path

# 结合量化和剪枝的优化流程
def optimize_model_for_deployment(encoder_path, decoder_path):
    # 加载模型
    encoder = Encoder().to(device)
    decoder = Decoder().to(device)
    encoder.load_state_dict(torch.load(encoder_path, map_location=device))
    decoder.load_state_dict(torch.load(decoder_path, map_location=device))
    
    # 1. 剪枝
    # 对编码器进行剪枝
    parameters_to_prune_encoder = []
    for name, module in encoder.named_modules():
        if isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d):
            parameters_to_prune_encoder.append((module, 'weight'))
    
    # 全局剪枝：移除20%的权重
    prune.global_unstructured(
        parameters_to_prune_encoder,
        pruning_method=prune.L1Unstructured,
        amount=0.2,
    )
    
    # 使剪枝永久化
    for module, name in parameters_to_prune_encoder:
        prune.remove(module, name)
    
    # 对解码器进行类似的剪枝...
    
    # 2. 量化
    # 将模型移至CPU
    encoder_cpu = encoder.cpu()
    decoder_cpu = decoder.cpu()
    
    # 准备示例输入
    example_cover = torch.randn(1, 3, 256, 256)
    example_secret = torch.randn(1, 3, 256, 256)
    
    # 为编码器创建包装器
    class EncoderWrapper(nn.Module):
        def __init__(self, encoder):
            super(EncoderWrapper, self).__init__()
            self.encoder = encoder
        
        def forward(self, cover, secret):
            return self.encoder(cover, secret)
    
    # 动态量化编码器包装器
    encoder_wrapper = EncoderWrapper(encoder_cpu)
    traced_encoder = torch.jit.trace(encoder_wrapper, (example_cover, example_secret))
    quantized_encoder = torch.quantization.quantize_dynamic(
        traced_encoder,
        {nn.Linear, nn.Conv2d, nn.ConvTranspose2d},
        dtype=torch.qint8
    )
    
    # 动态量化解码器
    traced_decoder = torch.jit.trace(decoder_cpu, example_cover)
    quantized_decoder = torch.quantization.quantize_dynamic(
        traced_decoder,
        {nn.Linear, nn.Conv2d, nn.ConvTranspose2d},
        dtype=torch.qint8
    )
    
    # 保存优化后的模型
    torch.jit.save(quantized_encoder, 'encoder_optimized.pt')
    torch.jit.save(quantized_decoder, 'decoder_optimized.pt')
    
    print("模型优化完成并保存")
    return 'encoder_optimized.pt', 'decoder_optimized.pt'

# 使用TensorRT加速模型推理
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 初始化TensorRT日志记录器
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def onnx_to_tensorrt(onnx_file_path, engine_file_path, precision='fp16'):
    """将ONNX模型转换为TensorRT引擎"""
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    # 解析ONNX文件
    with open(onnx_file_path, 'rb') as model:
        parser.parse(model.read())
    
    # 配置构建器
    config = builder.create_builder_config()
    if precision == 'fp16' and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    elif precision == 'int8' and builder.platform_has_fast_int8:
        config.set_flag(trt.BuilderFlag.INT8)
        # 这里需要提供校准数据，简化示例中省略
    
    # 设置最大批处理大小和工作空间大小
    builder.max_batch_size = 1
    config.max_workspace_size = 1 << 30  # 1GB
    
    # 构建引擎
    serialized_engine = builder.build_serialized_network(network, config)
    
    # 保存引擎
    with open(engine_file_path, 'wb') as f:
        f.write(serialized_engine)
    
    print(f'TensorRT engine saved to {engine_file_path}')
    return engine_file_path

# 将ONNX模型转换为TensorRT引擎
onnx_to_tensorrt('encoder.onnx', 'encoder_tensorrt.engine', precision='fp16')
onnx_to_tensorrt('decoder.onnx', 'decoder_tensorrt.engine', precision='fp16')

# 创建TensorRT推理上下文类
class TensorRTInferencer:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            runtime = trt.Runtime(self.logger)
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        
        # 分配设备内存
        self.inputs = []
        self.outputs = []
        self.bindings = []
        self.stream = cuda.Stream()
        
        for binding in range(self.engine.num_bindings):
            size = trt.volume(self.engine.get_binding_shape(binding)) * self.engine.max_batch_size
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            if self.engine.binding_is_input(binding):
                self.inputs.append({'host': host_mem, 'device': device_mem})
            else:
                self.outputs.append({'host': host_mem, 'device': device_mem})
    
    def infer(self, input_data):
        # 将输入数据复制到主机内存
        np.copyto(self.inputs[0]['host'], input_data.ravel())
        
        # 异步复制数据到设备
        for inp in self.inputs:
            cuda.memcpy_htod_async(inp['device'], inp['host'], self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 异步复制结果到主机
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
        
        # 同步流
        self.stream.synchronize()
        
        # 返回输出
        return self.outputs[0]['host']
    
    def __del__(self):
        # 释放资源
        del self.stream
        for inp in self.inputs:
            del inp['device']
        for out in self.outputs:
            del out['device']

# 使用TensorRT进行推理的示例
def tensorrt_steganography(encoder_engine_path, decoder_engine_path, cover_image, secret_image):
    # 创建推理器
    encoder_inferencer = TensorRTInferencer(encoder_engine_path)
    decoder_inferencer = TensorRTInferencer(decoder_engine_path)
    
    # 预处理图像
    cover_tensor = preprocess_pil_image(cover_image)
    secret_tensor = preprocess_pil_image(secret_image)
    
    # 合并输入（如果编码器的ONNX模型使用合并输入）
    combined_input = torch.cat([cover_tensor, secret_tensor], dim=1)
    
    # 执行编码
    stego_data = encoder_inferencer.infer(combined_input.cpu().numpy())
    
    # 重塑结果
    stego_tensor = torch.from_numpy(stego_data).reshape(1, 3, 256, 256).to(device)
    
    # 执行解码
    secret_data = decoder_inferencer.infer(stego_tensor.cpu().numpy())
    decoded_tensor = torch.from_numpy(secret_data).reshape(1, 3, 256, 256).to(device)
    
    # 后处理
    stego_image = postprocess_tensor(stego_tensor, cover_image.size)
    decoded_image = postprocess_tensor(decoded_tensor, secret_image.size)
    
    return stego_image, decoded_image

# 多线程与批处理优化示例
from concurrent.futures import ThreadPoolExecutor
import queue

def batch_process(images, batch_size=8, max_workers=4):
    """批量处理图像"""
    results = []
    
    # 创建线程池
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 分批次提交任务
        futures = []
        for i in range(0, len(images), batch_size):
            batch = images[i:i+batch_size]
            future = executor.submit(process_batch, batch)
            futures.append(future)
        
        # 收集结果
        for future in futures:
            results.extend(future.result())
    
    return results

def process_batch(batch):
    """处理单个批次的图像"""
    batch_results = []
    
    # 准备批次数据
    batch_tensors = []
    for image_path in batch:
        image = Image.open(image_path).convert('RGB')
        image_tensor = preprocess_pil_image(image)
        batch_tensors.append(image_tensor)
    
    # 合并为批次
    batch_tensor = torch.cat(batch_tensors, dim=0)
    
    # 处理批次（这里以解码为例）
    with torch.no_grad():
        outputs = decoder(batch_tensor)
    
    # 处理结果
    for i, output in enumerate(outputs):
        result = postprocess_tensor(output.unsqueeze(0))
        batch_results.append(result)
    
    return batch_results

# 队列处理示例
class SteganographyProcessor:
    def __init__(self, encoder, decoder, max_workers=4):
        self.encoder = encoder
        self.decoder = decoder
        self.input_queue = queue.Queue()
        self.output_queue = queue.Queue()
        self.max_workers = max_workers
        self.running = False
    
    def start(self):
        """启动处理线程"""
        self.running = True
        self.workers = []
        
        for _ in range(self.max_workers):
            worker = Thread(target=self._worker_thread)
            worker.daemon = True
            worker.start()
            self.workers.append(worker)
    
    def stop(self):
        """停止处理线程"""
        self.running = False
        for worker in self.workers:
            worker.join()
    
    def _worker_thread(self):
        """工作线程函数"""
        while self.running:
            try:
                # 从队列获取任务
                task = self.input_queue.get(timeout=1)
                
                # 处理任务
                if task['type'] == 'encode':
                    result = self._encode_task(task)
                elif task['type'] == 'decode':
                    result = self._decode_task(task)
                else:
                    result = {'status': 'error', 'message': 'Unknown task type'}
                
                # 将结果放入输出队列
                self.output_queue.put(result)
                
                # 标记任务完成
                self.input_queue.task_done()
            except queue.Empty:
                continue
            except Exception as e:
                self.output_queue.put({'status': 'error', 'message': str(e)})
    
    def _encode_task(self, task):
        """编码任务"""
        cover_image = Image.open(task['cover_path']).convert('RGB')
        secret_image = Image.open(task['secret_path']).convert('RGB')
        
        # 预处理
        cover_tensor = preprocess_pil_image(cover_image)
        secret_tensor = preprocess_pil_image(secret_image)
        
        # 执行编码
        with torch.no_grad():
            stego_tensor = self.encoder(cover_tensor, secret_tensor)
        
        # 后处理
        stego_image = postprocess_tensor(stego_tensor, cover_image.size)
        
        # 保存结果
        stego_image.save(task['output_path'])
        
        return {
            'status': 'success',
            'output_path': task['output_path']
        }
    
    def _decode_task(self, task):
        """解码任务"""
        stego_image = Image.open(task['stego_path']).convert('RGB')
        
        # 预处理
        stego_tensor = preprocess_pil_image(stego_image)
        
        # 执行解码
        with torch.no_grad():
            secret_tensor = self.decoder(stego_tensor)
        
        # 后处理
        secret_image = postprocess_tensor(secret_tensor, stego_image.size)
        
        # 保存结果
        secret_image.save(task['output_path'])
        
        return {
            'status': 'success',
            'output_path': task['output_path']
        }
    
    def submit_task(self, task):
        """提交任务"""
        self.input_queue.put(task)
    
    def get_result(self, timeout=None):
        """获取结果"""
        return self.output_queue.get(timeout=timeout)