OpenCV转pytorch

OpenCV的一些操纵转pytorch，从而有助于使用GPU加快，乃至导出onnx和转TensorRT
须要注意opencv的输入是numpy tensor，format是HW的2D张量大概HWC的3D张量，而pytorch一样平常是NCHW的4D大概CHW的3D张量。

Dilation腐蚀与膨胀

12: 腐蚀与膨胀 | 陌上见花开
https://blog.51cto.com/u_16175442/8629546

1. import cv2
2. import torch.nn.functional as F
5. def dilate_cv(img, dilate_factor=10):
6.     """
7.     input img is np 2D, HWC 3D
8.     """
9.     img = img.astype(np.uint8)
10.     img1 = cv2.dilate(
11.         img,
12.         np.ones((dilate_factor, dilate_factor), np.uint8),
13.         iterations=1
14.     )
15.     return img1
18. def dilate_torch(img, dilate_factor=10):
19.     """
20.     input img should be 3D CHW, or 4D NCHW
21.     """
22.     h, w = img.shape[-2:]
23.     img1 = F.max_pool2d(img, kernel_size=dilate_factor, stride=1, padding=dilate_factor//2)
24.     if dilate_factor % 2 == 0:
25.         img1 = img1[:, :, :h, :w]
26.     return img1

复制代码

贴一个DeepSeek转换的版本：

1. import torch
2. import numpy as np
4. def dilate_torch(img, dilate_factor=10):
5.     """
6.     input img is np 2D array
7.     """
8.     # 转换为PyTorch张量并添加batch和channel维度
9.     img_tensor = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float()
10.    
11.     # 创建最大池化层实现膨胀效果
12.     kernel_size = dilate_factor
13.     padding = dilate_factor // 2  # 保持输出尺寸与输入一致
14.     max_pool = torch.nn.MaxPool2d(kernel_size=kernel_size, stride=1, padding=padding)
15.    
16.     # 应用池化操作
17.     dilated_tensor = max_pool(img_tensor)
18.    
19.     # 转换回numpy数组并恢复原始维度
20.     dilated_img = dilated_tensor.squeeze().cpu().numpy()
21.    
22.     return dilated_img.astype(np.uint8)

复制代码

Resize

1. import cv2
2. from torchvision.transforms.functional import resize
3. from torchvision.transforms import InterpolationMode
5. img_cv = cv2.resize(img_hwc, (scale*W, scale*H), interpolation=cv2.INTER_NEAREST)
6. img_torch = resize(img_chw, (scale*H, scale*W), interpolation=InterpolationMode.NEAREST)

复制代码

须要注意的是opencv的resize和torch的resize效果不是完全对齐的，由于align方式的缘故原由。
颜色转换

1. bgr_cv = cv2.cvtColor(data_np, cv2.COLOR_RGB2BGR)
2. img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
4. def bgr2rgb_torch_nchw(bgr_nchw):
5.     b, g, r = bgr_nchw.split(split_size=1, dim=-3)
6.     rgb = torch.cat([r, g, b], dim=-3).numpy()
7.     return rgb
9. def rgb2bgr_torch_nchw(rgb_nchw):
10.     r, g, b = rgb_nchw.split(split_size=1, dim=-3)
11.     bgr = torch.cat([b, g, r], dim=-3)
12.     return bgr

复制代码

Blur

1. import torch
2. import numpy as np
3. import cv2
5. img_hwc = np.random.rand(*[256, 256, 3]).astype("float32")
6. img_chw = img_hwc.transpose([2, 0, 1])
7. img_chw_tc = torch.from_numpy(img_chw)
9. kernel_size = 3
11. img_blur_cv = cv2.blur(img_hwc, (kernel_size, kernel_size))
12. img_blur_cv_chw = img_blur_cv.transpose([2, 0, 1])
14. def mean_blur_torch(img_chw, kernel_size):
15.     device = img_chw.device
16.     dtype = img_chw.dtype
18.     pad_l = kernel_size // 2
19.     pad_r = kernel_size // 2
20.     if kernel_size % 2 == 0:
21.         pad_r = pad_r-1
23.     img_chw1 = torch.nn.functional.pad(img_chw, pad=[pad_l, pad_r, pad_l, pad_r], mode='reflect')
24.     weight = torch.ones(*(3, 1, kernel_size, kernel_size), dtype=dtype, device=device)/kernel_size/kernel_size
25.     img_blur_chw = torch.nn.functional.conv2d(img_chw1, weight, padding=0, groups=3)
26.     return img_blur_chw
28. img_blur_torch_chw = mean_blur_torch(img_chw_tc, kernel_size)
29. img_blur_torch_chw = img_blur_torch_chw.numpy()
31. error = np.abs(img_blur_cv_chw - img_blur_torch_chw)
32. print("error:", np.max(error), np.mean(error))

复制代码

1. def to_tensor_torch(tensor):
2.     # use torch tensor but not numpy as input
3.     # hwc to chw / nhwc to nchw transpose, dtype conversion, rescale
4.     if len(tensor.shape) == 2:
5.         tensor = tensor.unsqueeze(dim=-1)
7.     # transpose
8.     if len(tensor.shape) == 3:
9.         tensor = tensor.permute([2, 0, 1])
10.     elif len(tensor.shape) == 4:
11.         tensor = tensor.permute([0, 3, 1, 2])
12.     else:
13.         raise ValueError("unsupported")
15.     is_uint8 = tensor.dtype == torch.uint8
16.     # dtype conversion
17.     tensor = tensor.to(torch.float32)
18.     # rescale
19.     if is_uint8:
20.         tensor = tensor / 255.0
21.     return tensor

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！qidao123.com:ToB企服之家，中国第一个企服评测及软件市场,开放入驻,技术点评得现金