使用Pytoch实现Opencv warpAffine方法

随着深度学习的不断发展，GPU/NPU的算力也越来越强，对于一些传统CV计算也希望能够直接在GPU/NPU上进行，例如Opencv的warpAffine方法。Opencv的warpAffine的功能主要是做仿射变换，如果不了解仿射变换的请自行了解。由于Pytorch的图像坐标系（图像左上角对应坐标(-1, -1)右下角对应坐标(1, 1)）与Opencv的坐标系（图像左上角对应坐标(0, 0)右下角对应坐标(w - 1, h - 1)）有差异，故无法直接使用Opencv的warp矩阵对Pytorch数据进行变换。
主要参考文章：https://zhuanlan.zhihu.com/p/349741938

本文逻辑推理部分主要是参照上述的参考文章，这里再简单推导一遍。后面会给出基于该公式推导的Pytorch实现。

下面公式简单介绍了原始图片中 $x_1, y_1)$ 点通过仿射变化到输出图片 $x_2, y_2)$ 点的过程，假设 $(x, y)$ 对应Opencv图像坐标系。

$\begin{bmatrix} x_2\\ y_2 \\ 1 \end{bmatrix} = \begin{bmatrix} a & b & c\\ d & e & f\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} x_1\\ y_1 \\ 1 \end{bmatrix}$
现在要将Opencv图像坐标系下的 $x_1, y_1)$ 点映射到Pytorch的图像坐标系下 $u_1, v_1)$ 点，由于Pytorch的图像坐标系是从-1到1，所以对Opencv的坐标做如下变化即可。注，由于Opencv坐标从0开始，所以对于原图宽为src_w，高为src_h实际右下角的坐标应该是 $src_w - 1, src_h - 1)$ 。
$u_1 = \frac{x_1 - \frac{src_w - 1}{2} }{\frac{src_w - 1}{2}} = \frac{2x_1}{src_w - 1} -1$
$v_1 = \frac{y_1 - \frac{src_h - 1}{2} }{\frac{src_h - 1}{2}} = \frac{2y_1}{src_h - 1} -1$
写成矩阵乘的形式：
$\begin{bmatrix} u_1\\ v_1 \\ 1 \end{bmatrix} = \begin{bmatrix} \frac{2}{src_w - 1} & 0 & -1\\ 0 & \frac{2}{src_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} x_1\\ y_1 \\ 1 \end{bmatrix}$

那么同理将仿射变化后Opencv图像坐标系下的 $x_2, y_2)$ 点映射到Pytorch的图像坐标系下 $u_2, v_2)$ 点，其中dst_w为仿射变化后输出图片的宽度，dst_h为仿射变化后输出图片的高度：
$\begin{bmatrix} u_2\\ v_2 \\ 1 \end{bmatrix} = \begin{bmatrix} \frac{2}{dst_w - 1} & 0 & -1\\ 0 & \frac{2}{dst_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} x_2\\ y_2 \\ 1 \end{bmatrix}$
然后将上面两个公式代入最开始的仿射变化公式中：
$\begin{bmatrix} \frac{2}{dst_w - 1} & 0 & -1\\ 0 & \frac{2}{dst_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix}^{-1} \begin{bmatrix} u_2\\ v_2 \\ 1 \end{bmatrix} = \begin{bmatrix} a & b & c\\ d & e & f\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} \frac{2}{src_w - 1} & 0 & -1\\ 0 & \frac{2}{src_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix}^{-1} \begin{bmatrix} u_1\\ v_1 \\ 1 \end{bmatrix}$
整理得到：
$\begin{bmatrix} u_2\\ v_2 \\ 1 \end{bmatrix} = \begin{bmatrix} \frac{2}{dst_w - 1} & 0 & -1\\ 0 & \frac{2}{dst_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} a & b & c\\ d & e & f\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} \frac{2}{src_w - 1} & 0 & -1\\ 0 & \frac{2}{src_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix}^{-1} \begin{bmatrix} u_1\\ v_1 \\ 1 \end{bmatrix}$
引用参考文章中大佬的原话，这个暂时没在Pytorch官方文档中找到，但是通过实验，确实如此。

affine_grid定义为目标图到原图的变换

所以，Pytorch中使用的theta实际是从 $u_2, v_2)$ 到 $u_1, v_1)$ 的矩阵：

$\begin{bmatrix} u_1\\ v_1 \\ 1 \end{bmatrix} = \begin{bmatrix} \frac{2}{src_w - 1} & 0 & -1\\ 0 & \frac{2}{src_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} a & b & c\\ d & e & f\\ 0 & 0 & 1 \end{bmatrix}^{-1} \begin{bmatrix} \frac{2}{dst_w - 1} & 0 & -1\\ 0 & \frac{2}{dst_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix}^{-1} \begin{bmatrix} u_2\\ v_2 \\ 1 \end{bmatrix}$
故Opencv使用的theta到Pytorch的theta变换过程如下：
$theta_{(pytorch)} = \begin{bmatrix} \frac{2}{src_w - 1} & 0 & -1\\ 0 & \frac{2}{src_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix} {theta}^{-1}_{(opencv)} \begin{bmatrix} \frac{2}{dst_w - 1} & 0 & -1\\ 0 & \frac{2}{dst_h - 1} & -1\\ 0 & 0 & 1 \end{bmatrix}^{-1}$

最后给出对应代码实现：

"""
pip install numpy
pip install opencv-python
pip install opencv-python-headless
"""
import numpy as np
import cv2
import torch
import torch.nn.functional as F


def cal_torch_theta(opencv_theta: np.ndarray, src_h: int, src_w: int, dst_h: int, dst_w: int):
    m = np.concatenate([opencv_theta, np.array([[0., 0., 1.]], dtype=np.float32)])
    m_inv = np.linalg.inv(m)

    a = np.array([[2 / (src_w - 1), 0., -1.],
                  [0., 2 / (src_h - 1), -1.],
                  [0., 0., 1.]], dtype=np.float32)

    b = np.array([[2 / (dst_w - 1), 0., -1.],
                  [0., 2 / (dst_h - 1), -1.],
                  [0., 0., 1.]], dtype=np.float32)
    b_inv = np.linalg.inv(b)

    pytorch_m = a @ m_inv @ b_inv
    return torch.as_tensor(pytorch_m[:2], dtype=torch.float32)


def main():
    img_bgr = cv2.imread("1.png")
    src_h, src_w, _ = img_bgr.shape
    print(f"src image h:{src_h}, w:{src_w}")
    dst_h = src_h * 2
    dst_w = src_w * 2
    print(f"dst image h:{src_h}, w:{src_w}")

    theta = cv2.getRotationMatrix2D(center=(src_w // 2, src_h // 2), angle=-30, scale=2)
    # using opencv warpAffine
    warp_img_bgr = cv2.warpAffine(src=img_bgr,
                                  M=theta,
                                  dsize=(dst_w, dst_h),
                                  flags=cv2.INTER_LINEAR,
                                  borderValue=(0, 0, 0))
    cv2.imwrite("warp_img.jpg", warp_img_bgr)

    # using pytorch grid_sample
    torch_img_bgr = torch.as_tensor(img_bgr, dtype=torch.float32).unsqueeze(0).permute([0, 3, 1, 2])  # [N,C,H,W]
    torch_theta = cal_torch_theta(theta, src_h, src_w, dst_h, dst_w).unsqueeze(0)  # [N, 2, 3]
    grid = F.affine_grid(torch_theta, size=[1, 3, dst_h, dst_w])
    torch_warp_img_bgr = F.grid_sample(torch_img_bgr, grid=grid, mode="bilinear", padding_mode="zeros")

    torch_warp_img_bgr = torch_warp_img_bgr.permute([0, 2, 3, 1]).squeeze(0)  # [H, W, C]
    cv2.imwrite("torch_warp_img.jpg", torch_warp_img_bgr.numpy())

    # save concat img
    cv2.imwrite("compare_warp_img.jpg",
                np.concatenate([warp_img_bgr, torch_warp_img_bgr.numpy()], axis=1))


if __name__ == '__main__':
    main()