图像畸变矫正算法-STN

论文题目：Spatial Transformer Networks

STN.png

• 定位网络：一般是由一个全连接网络实现，该网络最终输出6个数值，将其转换为2x3矩阵，其中每两个为一组，分别表示仿射变换中的平移、旋转、缩放所对应的参数。
• 网格生成器：创建一个与输出图片大小相同的矩阵，把该矩阵与仿射参数矩阵相乘将得到的结果作为图片中每个像素点对应于原图的坐标。
• 采样器：使用坐标映射中每个像素点的坐标值，在原始图片上取相应的像素，并将其填充到目标图片中，得到整幅目标图片。
  通俗一点讲就是通过训练学习一个仿射变换矩阵，而不是人为定义这个矩阵，这个仿射变换矩阵可以将形变的图像大致恢复回来。
  注意：STN更适合用于分类或文字识别这样输出结果跟图像位置不太相关的任务，想目标检测或语义分割任务的改动量就比较大，由于STN对输入图像做了形变所以标签也要做相应变换的。
  代码如下所示：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

        # Spatial transformer localization-network
        self.localization = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True)
        )

        # Regressor for the 3 * 2 affine matrix
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 3 * 3, 32),
            nn.ReLU(True),
            nn.Linear(32, 3 * 2)
        )

        # Initialize the weights/bias with identity transformation
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))

    # Spatial transformer network forward function
    def stn(self, x):
        xs = self.localization(x)
        # xs = xs.view(-1, 10 * 3 * 3)
        xs = torch.flatten(xs, 1)
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)

        return x

    def forward(self, x):
        # transform the input
        x = self.stn(x)

        # Perform the usual forward pass
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

这是一个分类网络中在开头部分增加了STN模块，用于消除或降低图像形变的影响，从而降低后续模型的分类压力