使用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_1, y_1) (x1,y1)点通过仿射变化到输出图片 ( x 2 , y 2 ) (x_2, y_2) (x2,y2)点的过程，假设 ( x , y ) (x, y) (x,y)对应Opencv图像坐标系。

[ x 2 y 2 1 ] = [ a b c d e f 0 0 1 ] [ x 1 y 1 1 ] \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} x2y21 = ad0be0cf1 x1y11

现在要将Opencv图像坐标系下的 ( x 1 , y 1 ) (x_1, y_1) (x1,y1)点映射到Pytorch的图像坐标系下 ( u 1 , v 1 ) (u_1, v_1) (u1,v1)点，由于Pytorch的图像坐标系是从-1到1，所以对Opencv的坐标做如下变化即可。注，由于Opencv坐标从0开始，所以对于原图宽为src_w，高为src_h实际右下角的坐标应该是 ( s r c w − 1 , s r c h − 1 ) (src_w - 1, src_h - 1) (srcw−1,srch−1)。
u 1 = x 1 − s r c w − 1 2 s r c w − 1 2 = 2 x 1 s r c w − 1 − 1 u_1 = \frac{x_1 - \frac{src_w - 1}{2} }{\frac{src_w - 1}{2}} = \frac{2x_1}{src_w - 1} -1 u1=2srcw−1x1−2srcw−1=srcw−12x1−1
v 1 = y 1 − s r c h − 1 2 s r c h − 1 2 = 2 y 1 s r c h − 1 − 1 v_1 = \frac{y_1 - \frac{src_h - 1}{2} }{\frac{src_h - 1}{2}} = \frac{2y_1}{src_h - 1} -1 v1=2srch−1y1−2srch−1=srch−12y1−1

写成矩阵乘的形式：
[ u 1 v 1 1 ] = [ 2 s r c w − 1 0 − 1 0 2 s r c h − 1 − 1 0 0 1 ] [ x 1 y 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} u1v11 = srcw−12000srch−120−1−11 x1y11

那么同理将仿射变化后Opencv图像坐标系下的 ( x 2 , y 2 ) (x_2, y_2) (x2,y2)点映射到Pytorch的图像坐标系下 ( u 2 , v 2 ) (u_2, v_2) (u2,v2)点，其中dst_w为仿射变化后输出图片的宽度，dst_h为仿射变化后输出图片的高度：
[ u 2 v 2 1 ] = [ 2 d s t w − 1 0 − 1 0 2 d s t h − 1 − 1 0 0 1 ] [ x 2 y 2 1 ] \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} u2v21 = dstw−12000dsth−120−1−11 x2y21

然后将上面两个公式代入最开始的仿射变化公式中：
[ 2 d s t w − 1 0 − 1 0 2 d s t h − 1 − 1 0 0 1 ] − 1 [ u 2 v 2 1 ] = [ a b c d e f 0 0 1 ] [ 2 s r c w − 1 0 − 1 0 2 s r c h − 1 − 1 0 0 1 ] − 1 [ u 1 v 1 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} = \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} dstw−12000dsth−120−1−11 −1 u2v21 = ad0be0cf1 srcw−12000srch−120−1−11 −1 u1v11

整理得到：
[ u 2 v 2 1 ] = [ 2 d s t w − 1 0 − 1 0 2 d s t h − 1 − 1 0 0 1 ] [ a b c d e f 0 0 1 ] [ 2 s r c w − 1 0 − 1 0 2 s r c h − 1 − 1 0 0 1 ] − 1 [ u 1 v 1 1 ] \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} u2v21 = dstw−12000dsth−120−1−11 ad0be0cf1 srcw−12000srch−120−1−11 −1 u1v11

引用参考文章中大佬的原话，这个暂时没在Pytorch官方文档中找到，但是通过实验，确实如此。

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

所以，Pytorch中使用的theta实际是从 ( u 2 , v 2 ) (u_2, v_2) (u2,v2)到 ( u 1 , v 1 ) (u_1, v_1) (u1,v1)的矩阵：

[ u 1 v 1 1 ] = [ 2 s r c w − 1 0 − 1 0 2 s r c h − 1 − 1 0 0 1 ] [ a b c d e f 0 0 1 ] − 1 [ 2 d s t w − 1 0 − 1 0 2 d s t h − 1 − 1 0 0 1 ] − 1 [ u 2 v 2 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} u1v11 = srcw−12000srch−120−1−11 ad0be0cf1 −1 dstw−12000dsth−120−1−11 −1 u2v21

故Opencv使用的theta到Pytorch的theta变换过程如下：
t h e t a ( p y t o r c h ) = [ 2 s r c w − 1 0 − 1 0 2 s r c h − 1 − 1 0 0 1 ] t h e t a ( o p e n c v ) − 1 [ 2 d s t w − 1 0 − 1 0 2 d s t h − 1 − 1 0 0 1 ] − 1 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} theta(pytorch)= srcw−12000srch−120−1−11 theta(opencv)−1 dstw−12000dsth−120−1−11 −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得到的图片。可以看到基本是一致，如果使用专业的图像对比工具还是能看到像素差异（很难完全对齐）。