使用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图像坐标系。

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

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

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

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

$$\left[\begin{array}{c}{u}_1\\ v_1\\ 1\end{array}\right]=\left[\begin{array}{ccc}\frac{2}{sr{c}_w-1}& 0& -1\\ 0& \frac{2}{sr{c}_h-1}& -1\\ 0& 0& 1\end{array}\right]{\left[\begin{array}{ccc}a& b& c\\ d& e& f\\ 0& 0& 1\end{array}\right]}^{-1}{\left[\begin{array}{ccc}\frac{2}{ds{t}_w-1}& 0& -1\\ 0& \frac{2}{ds{t}_h-1}& -1\\ 0& 0& 1\end{array}\right]}^{-1}\left[\begin{array}{c}{u}_2\\ v_2\\ 1\end{array}\right]$$
故Opencv使用的theta到Pytorch的theta变换过程如下：
$$thet{a}_{(pytorch)}=\left[\begin{array}{ccc}\frac{2}{sr{c}_w-1}& 0& -1\\ 0& \frac{2}{sr{c}_h-1}& -1\\ 0& 0& 1\end{array}\right]{theta}_{(opencv)}^{-1}{\left[\begin{array}{ccc}\frac{2}{ds{t}_w-1}& 0& -1\\ 0& \frac{2}{ds{t}_h-1}& -1\\ 0& 0& 1\end{array}\right]}^{-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()

下图是生成的compare_warp_img.jpg图片，左边是通过Opencv warpAffine得到的图片，右边是通过Pytorch grid_sample得到的图片。可以看到基本是一致，如果使用专业的图像对比工具还是能看到像素差异（很难完全对齐）。