tps_stn_pytorch

基于薄板样条（TPS）的空间变换网络（STN）的 PyTorch 实现。

简介

空间变换网络（STN）是由 DeepMind 提出的一种强大的神经网络架构 [1]。STN 通过在输入图像被送入常规分类网络之前自动对其进行校正，从而实现真正的空间不变性。STN 最令人惊叹之处在于它是端到端可微分的，并且可以直接插入现有的网络架构中（如 AlexNet、ResNet 等），无需任何额外的监督信号。

原始 STN 论文 [1] 对三种特定的变换形式进行了实验：仿射变换、透视变换和 薄板样条变换（TPS）。其中我认为 TPS 是最强大的变换方式，因为它可以以任意方式扭曲图像。如下所示，我可以将我的头像进行扭曲：

变为

TPS-STN 已被用于 OCR 应用程序 [2]。在这篇论文中，TPS-STN 用于在将扭曲的文本图像送入常规 OCR 文本识别模型之前，自动对其进行校正：

依赖项

• Python3
• PyTorch
• Torchvision
• Numpy
• Matplotlib
• Pillow / PIL
• imageio

我使用 imageio 来创建 GIF 可视化效果。只需运行 pip install imageio 即可安装。

运行

python mnist_train.py --model unbounded_stn --angle 90 --grid_size 4
python mnist_visualize.py --model unbounded_stn --angle 90 --grid_size 4
python mnist_make_gif.py --model unbounded_stn --angle 90 --grid_size 4

随后，PNG 和 GIF 结果将分别保存在 ./image/unbounded_stn_angle60_grid4/ 和 ./gif/unbounded_stn_angle60_grid4/ 目录中。

您可以尝试其他模型架构、MNIST 数据集随机旋转角度以及 TPS 网格大小的不同组合。详细信息如下。

参数

共有三个可控参数：--model、--angle 和 --grid_size。

--model: 字符串，必填

• 使用 no_stn 时，STN 模块将被丢弃，仅保留单一的 CNN 分类器。
• 使用 bounded_stn 时，定位网络的输出会被 F.tanh 压缩到 [-1, 1] 范围内，正如 [2] 中所做的那样。
• 使用 unbounded_stn 时，定位网络的输出不会被压缩。

--angle: 整数，默认为 60

• MNIST 数据集中的样本将被随机旋转 [-angle, angle] 范围内的角度。

--grid_size: 整数，默认为 4

• 使用 (grid_size x grid_size) 个控制点来定义薄板样条变换。

测试准确率曲线

可视化 TPS-STN

当 angle = 90 时，结果通常较差：

当使用 bounded_stn 且 grid_size <= 3 时，结果也不理想：

但若 grid_size >= 4，则效果较好：

而使用 unbounded_stn 时，结果则较为理想：

当然，在每种组合下都可能存在表现不佳的情况。您可以通过 百度网盘 下载我所有的 GIF 文件（文件大小为 2GB）。

参考文献

[1] Spatial Transformer Networks

[2] Robust Scene Text Recognition with Automatic Rectification

[3] 数值方法——薄板样条插值（Thin-Plate Spline）

tps_stn_pytorch 快速上手指南

tps_stn_pytorch 是一个基于 PyTorch 实现的带有薄板样条（Thin Plate Spline, TPS）的空间变换网络（STN）。该工具能够自动校正输入图像的几何形变，无需额外监督信号即可端到端地嵌入现有网络架构（如 ResNet、AlexNet），特别适用于 OCR 文本校正等场景。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：Python 3.x
• 核心依赖：
  • PyTorch
  • Torchvision
  • Numpy
  • Matplotlib
  • Pillow (PIL)
  • imageio (用于生成 GIF 可视化)

国内加速建议： 安装 PyTorch 时，推荐使用清华或中科大镜像源以提升下载速度。例如： pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆项目代码

   git clone https://github.com/warbean/tps_stn_pytorch.git
   cd tps_stn_pytorch
   

2. 安装 Python 依赖包 建议使用国内镜像源安装剩余依赖：

   pip install numpy matplotlib pillow imageio -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：PyTorch 和 Torchvision 需根据您的 CUDA 版本单独安装，若未安装请先执行上述加速安装命令)

基本使用

本项目提供了针对 MNIST 数据集的完整训练、可视化和 GIF 生成脚本。以下是最简单的运行示例，演示如何训练一个无边界 TPS-STN 模型，并处理旋转角度为 90 度、网格大小为 4 的图像。

1. 训练模型

运行训练脚本，模型将自动学习如何校正随机旋转的 MNIST 数字图像。

python mnist_train.py --model unbounded_stn --angle 90 --grid_size 4

2. 可视化结果

训练完成后，生成静态 PNG 图像以观察校正效果。

python mnist_visualize.py --model unbounded_stn --angle 90 --grid_size 4

3. 生成动态演示 (GIF)

生成展示变换过程的 GIF 动图。

python mnist_make_gif.py --model unbounded_stn --angle 90 --grid_size 4

输出说明： 执行上述命令后，生成的图片将保存在以下目录：

• 静态图片：./image/unbounded_stn_angle60_grid4/
• 动态 GIF：./gif/unbounded_stn_angle60_grid4/ (注：目录名称中的角度值可能根据实际运行参数略有不同)

参数说明

您可以根据需求调整以下三个核心参数：

• --model: 选择模型类型
  • no_stn: 不使用 STN，仅保留普通 CNN 分类器（基线对比）。
  • bounded_stn: 定位网络输出通过 tanh 限制在 [-1, 1] 区间。
  • unbounded_stn: 定位网络输出不进行限制（推荐尝试）。
• --angle: 设置 MNIST 样本的随机旋转角度范围 [-angle, angle]，默认为 60。
• --grid_size: 定义 TPS 变换的控制点网格大小 (grid_size x grid_size)，默认为 4。建议 grid_size >= 4 以获得较好的校正效果。