【idea】deepOnet在图像畸变矫正方面的潜力研究

2025-04-25

deepOnet

Word count: 3.6k | Reading time≈ 14 min

之前我们已经讲过了deepOnet到底是怎么一回事，大家也可以从中看出这个网络在物理规律拟合上的潜力。但是我觉得这个网络的潜力远远不止物理规律拟合这一个方面。

deepOnet网络做的，从宏观来看，绝对不止是局限于学习函数的映射。他学习的是如何根据一种变化趋势来得到另外一种变化趋势。放在我们原来的“算子拟合”上去理解，就是其学会的了如何根据一种变化方式，来推理出同样的事物的另一种变化方式。

如果我们将其用在图像处理上，或者更加精确一些——图像畸变矫正上的？

技术背景

图像畸变是指由于相机镜头的物理特性或图像采集过程中的误差，导致图像中的物体形状发生扭曲的现象。在光学成像系统中，光线通过透镜时会发生折射，这种折射并非完全理想，从而导致图像产生畸变。常见的图像畸变类型包括桶形畸变、枕形畸变和复杂畸变。

桶形畸变表现为图像边缘向外弯曲，其数学模型可以表示为：

\begin{cases} x' = x(1 + k_1r^2 + k_2r^4 + \dots) \\ y' = y(1 + k_1r^2 + k_2r^4 + \dots) \end{cases}

其中， $(x, y)$ 是原始图像中的坐标， $(x', y')$ 是畸变后的坐标， $r^2 = x^2 + y^2$ ， $k_1$ 和 $k_2$ 是畸变系数。当 $k_1 > 0$ 时，表现为桶形畸变。

枕形畸变则表现为图像边缘向内收缩，其数学模型与桶形畸变类似，但畸变系数 $k_1 < 0$ ：

\begin{cases} x' = x(1 + k_1r^2 + k_2r^4 + \dots) \\ y' = y(1 + k_1r^2 + k_2r^4 + \dots) \end{cases}

复杂畸变则是由多种因素混合导致的非线性畸变，难以用简单的数学模型描述，通常需要更复杂的函数来建模。

传统的图像畸变校正方法主要基于物理模型，如Brown-Conrady模型、Tsai模型等。这些方法通常需要进行相机标定，以估计相机的内参和畸变系数。虽然传统方法在某些场景下能够取得较好的效果，但它们存在一些局限性。例如，传统方法通常需要手动选择特征点，标定过程繁琐且耗时；对于复杂的非线性畸变，传统模型的表达能力有限；此外，传统方法难以适应不同类型的相机和畸变场景。

深度学习技术的发展为图像畸变校正提供了新的解决方案。与传统方法相比，深度学习方法具有以下优势：首先，深度学习方法可以通过大量数据自动学习畸变模式和校正映射，无需手动设计复杂的数学模型；其次，深度学习方法具有较强的泛化能力，能够处理多种类型的畸变；此外，深度学习方法可以实现端到端的训练和推理，提高了处理效率和自动化程度。

系统设计与实现

整体架构

为了实现高效、准确的图像畸变校正，我设计了一个基于特征引导DeepONet的图像畸变校正系统。该系统主要由四个核心模块组成：增强型畸变数据集、特征引导的DeepONet模型、高级训练框架和评估与可视化模块。

增强型畸变数据集模块负责生成包含多种畸变类型的训练数据。该模块基于CIFAR-10数据集，通过应用不同的畸变模型和数据增强技术，生成多样化的畸变图像样本。特征引导的DeepONet模型是系统的核心，它通过分离特征提取和坐标处理，学习从畸变图像到无畸变图像的映射关系。高级训练框架采用混合精度训练和多损失函数优化策略，提高了模型的训练效率和性能。评估与可视化模块则用于定量评估模型性能并可视化校正结果，帮助用户直观了解模型的效果。

数据生成模块

数据是深度学习模型训练的基础。为了训练一个能够处理多种畸变类型的模型，我创建了一个增强型畸变数据集。该数据集基于CIFAR-10数据集，通过应用三种不同的畸变模型（桶形畸变、枕形畸变和复杂畸变）生成畸变图像。

在实现过程中，首先创建一个归一化的坐标网格，然后根据不同的畸变模型对网格进行变换。以下是具体实现代码：

def _apply_distortion(self, image, distortion_type):
    """应用物理畸变模型"""
    # 创建归一化网格
    x = torch.linspace(-1, 1, self.img_size)
    y = torch.linspace(-1, 1, self.img_size)
    grid_x, grid_y = torch.meshgrid(x, y)
    grid = torch.stack([grid_x, grid_y], dim=-1)

    # 桶形畸变模型
    if distortion_type == 'barrel':
        k = random.uniform(*self.distortion_params['barrel'])
        r = torch.norm(grid, dim=-1)
        scale = 1 + k * r**2
        distorted_grid = grid * scale.unsqueeze(-1)

    # 枕形畸变模型
    elif distortion_type == 'pincushion':
        k = random.uniform(*self.distortion_params['pincushion'])
        r = torch.norm(grid, dim=-1)
        scale = 1 - k * r**2
        distorted_grid = grid * scale.unsqueeze(-1)

    # 复杂畸变模型
    else:
        k = random.uniform(*self.distortion_params['complex'])
        complex_dist = torch.sin(3 * grid_x) * k + torch.cos(3 * grid_y) * k
        distorted_grid = grid + complex_dist.unsqueeze(-1)

    # 应用畸变
    distorted_image = F.grid_sample(
        image.unsqueeze(0),
        distorted_grid.unsqueeze(0),
        mode='bilinear',
        padding_mode='reflection',
        align_corners=False
    ).squeeze(0)

    return distorted_image

对于桶形畸变，使用公式 $scale = 1 + k \cdot r^2$ 对网格进行缩放，其中 $k$ 是随机生成的畸变系数， $r$ 是到图像中心的距离。对于枕形畸变，使用公式 $scale = 1 - k \cdot r^2$ 进行缩放。复杂畸变则使用更复杂的非线性变换，结合了正弦和余弦函数，能够模拟更复杂的畸变模式。

除了基本的畸变模型，数据生成模块还添加了随机噪声和模糊处理，进一步增强了数据的多样性。这些数据增强技术有助于提高模型的泛化能力，使其能够更好地应对实际应用中的各种畸变场景。

模型架构：特征引导的DeepONet

核心模型采用了DeepONet架构，并进行了针对性的改进，形成了特征引导的DeepONet模型。DeepONet是一种新型的神经网络架构，最初用于求解算子方程，近年来在科学计算和工程领域得到了广泛应用。在本项目中，我将DeepONet架构引入图像畸变校正领域，取得了良好的效果。

特征引导的DeepONet模型主要由四个部分组成：特征提取网络、分支网络、主干网络和偏移量预测网络。特征提取网络是一个卷积神经网络，用于提取输入畸变图像的全局特征。分支网络接收特征提取网络的输出，对特征进行进一步处理和转换。主干网络则处理空间坐标信息，将坐标映射到高维特征空间。偏移量预测网络将分支网络和主干网络的输出进行融合，预测每个像素的空间偏移量。

以下是模型的具体实现代码：

class FeatureGuidedDeepONet(nn.Module):
    def __init__(self, img_size=128):
        super().__init__()

        # 特征提取网络 - 获取全局畸变特征
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(3, 64, 5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 5, padding=2),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten()
        )

        # DeepONet分支网络 - 处理特征向量
        self.branch_net = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 256)
        )

        # DeepONet主干网络 - 处理坐标信息
        self.trunk_net = nn.Sequential(
            nn.Linear(2, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 256)
        )

        # 偏移量预测网络
        self.offset_predictor = nn.Sequential(
            nn.Linear(256 + 256, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 2)  # 预测x,y方向的偏移量
        )

        # 细节增强器
        self.detail_enhancer = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 3, 3, padding=1)
        )

    def forward(self, x):
        # 提取畸变特征
        features = self.feature_extractor(x)

        # 生成空间坐标网格
        batch_size, _, h, w = x.size()
        grid = self._create_grid(batch_size, h, w, x.device)
        grid_flat = grid.view(batch_size, h * w, 2)

        # 处理分支网络和主干网络
        branch_output = self.branch_net(features).unsqueeze(1).expand(-1, h * w, -1)
        trunk_output = self.trunk_net(grid_flat)

        # 结合特征和坐标信息，预测偏移量
        combined = torch.cat([branch_output, trunk_output], dim=-1)
        offset = self.offset_predictor(combined)
        offset_map = offset.view(batch_size, h, w, 2)

        # 应用偏移网格进行图像校正
        corrected_grid = grid + offset_map * 0.2
        output = F.grid_sample(
            x, corrected_grid,
            mode='bilinear',
            padding_mode='reflection',
            align_corners=False
        )

        # 细节增强
        output = output + self.detail_enhancer(output) * 0.2

        return output.clamp(0, 1)

模型的工作流程如下：首先，特征提取网络从输入的畸变图像中提取全局特征；然后，分支网络对这些特征进行处理，生成与畸变类型和强度相关的特征表示；同时，主干网络处理图像中每个像素的坐标信息；最后，偏移量预测网络将分支网络和主干网络的输出结合起来，预测每个像素的偏移量，从而实现图像的校正。

从数学角度来看，DeepONet的核心思想是将算子学习问题分解为两个子网络：分支网络 $g(\cdot)$ 和主干网络 $\phi(\cdot)$ 。对于输入函数 $u$ 和点 $x$ ，DeepONet的输出可以表示为：

Q(u)(x) \approx \sum_{k=1}^p g_k(u) \cdot \phi_k(x)

在图像畸变校正的场景中，分支网络 $g(\cdot)$ 对应于特征提取网络和分支网络的组合，用于提取畸变图像的特征；主干网络 $\phi(\cdot)$ 对应于主干网络，用于处理坐标信息；两者的组合通过偏移量预测网络实现，最终得到像素级的偏移量。

为了进一步提高模型的性能，我还添加了一个细节增强模块。该模块是一个小型的卷积神经网络，用于恢复图像的细节信息，提高校正后图像的质量。

训练与评估

为了训练特征引导的DeepONet模型，我设计了一个高级训练框架，采用了多项优化技术。首先，为了提高训练效率和减少显存占用，采用了混合精度训练技术。混合精度训练使用半精度浮点数（FP16）进行计算，同时保持关键步骤的单精度浮点数（FP32）精度，在不损失模型性能的前提下，显著提高了训练速度和内存利用率。

以下是训练框架的核心代码：

def train_model(model, train_loader, val_loader, optimizer, device, epochs=50, save_dir="results"):
    os.makedirs(save_dir, exist_ok=True)

    # 混合精度训练设置
    scaler = torch.cuda.amp.GradScaler() if device.type == 'cuda' else None

    # 损失函数
    reconstruction_loss = nn.L1Loss()
    detail_loss = nn.MSELoss()

    best_psnr = 0
    history = {'train_loss': [], 'val_psnr': [], 'val_ssim': []}

    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0
        train_bar = tqdm(train_loader, desc=f'Train Epoch {epoch+1}/{epochs}')

        for batch in train_bar:
            distorted = batch['distorted'].to(device)
            clean = batch['clean'].to(device)

            optimizer.zero_grad()

            # 混合精度前向传播
            with torch.cuda.amp.autocast(enabled=(device.type=='cuda')):
                corrected = model(distorted)
                loss = reconstruction_loss(corrected, clean)
                edge_loss = detail_loss(kornia.filters.sobel(corrected),
                                      kornia.filters.sobel(clean)) * 0.5
                total_loss = loss + edge_loss

            # 梯度缩放和反向传播
            if scaler:
                scaler.scale(total_loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                total_loss.backward()
                optimizer.step()

            train_loss += total_loss.item()
            train_bar.set_postfix(loss=f"{total_loss.item():.4f}")

        # 验证阶段
        model.eval()
        val_metrics = validate_model(model, val_loader, device)
        history['val_psnr'].append(val_metrics['psnr'])
        history['val_ssim'].append(val_metrics['ssim'])

        print(f"Epoch {epoch+1}/{epochs} | Train Loss: {train_loss/len(train_loader):.4f} | "
              f"Val PSNR: {val_metrics['psnr']:.2f} dB | Val SSIM: {val_metrics['ssim']:.4f}")

        # 保存最佳模型
        if val_metrics['psnr'] > best_psnr:
            best_psnr = val_metrics['psnr']
            torch.save(model.state_dict(), os.path.join(save_dir, 'best_model.pth'))

其次，为了平衡整体结构恢复和细节保留，设计了多损失函数组合策略。主要使用两种损失函数：L1重建损失和边缘细节损失。L1重建损失用于衡量校正图像与真实图像之间的像素级差异，促使模型恢复图像的整体结构：

L_{reconstruction} = \frac{1}{N} \sum_{i=1}^N |I_{corrected}(i) - I_{clean}(i)|

其中， $I_{corrected}$ 是校正后的图像， $I_{clean}$ 是真实无畸变的图像， $N$ 是图像中的像素总数。

边缘细节损失则使用Sobel算子提取图像的边缘信息，衡量校正图像与真实图像边缘信息的差异，帮助模型保留图像的细节特征：

L_{detail} = \frac{1}{N} \sum_{i=1}^N |\nabla I_{corrected}(i) - \nabla I_{clean}(i)|^2

其中， $\nabla I$ 表示使用Sobel算子计算的图像梯度。

最终的总损失函数是这两种损失的组合：

L_{total} = L_{reconstruction} + \lambda L_{detail}

在训练过程中，还采用了学习率自适应调整策略。使用ReduceLROnPlateau调度器，根据验证集性能自动调整学习率。当验证集性能在多个轮次内没有提升时，降低学习率，有助于模型跳出局部最优解，进一步提高性能。

为了评估模型的性能，使用了多种评估指标，包括峰值信噪比（PSNR）、结构相似性指数（SSIM）等。同时，设计了可视化模块，将校正结果与原始畸变图像和真实图像进行对比，直观展示模型的校正效果。那么最后其表现如何呢？我觉得我们可以直接看一下在极端扭曲的情况下模型的表现：

在极端的情况下，模型在结构和清晰度上进行了有力的还原。

Donate