使用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 Fdef 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_invreturn torch.as_tensor(pytorch_m[:2], dtype=torch.float32)def main():img_bgr = cv2.imread("1.png")src_h, src_w, _ = img_bgr.shapeprint(f"src image h:{src_h}, w:{src_w}")dst_h = src_h * 2dst_w = src_w * 2print(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 warpAffinewarp_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_sampletorch_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 imgcv2.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得到的图片。可以看到基本是一致，如果使用专业的图像对比工具还是能看到像素差异（很难完全对齐）。