发现一个宝藏回答:
代码
- 可以通过调节
kernel_size
实现腐蚀和膨胀的大小。
import cv2
import torch
import torch.nn.functional as F
import numpy as np
def dilate(bin_img, ksize=5):
# 腐蚀
src_size = bin_img.numpy().shape
pad = (ksize - 1) // 2
# bin_img = F.pad(bin_img, pad=[pad, pad, pad, pad], mode='reflect')
out = F.max_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0)
out = F.interpolate(out,
size=src_size[2:],
mode="bilinear")
return out
def erode(bin_img, ksize=5):
# 膨胀
out = 1 - dilate(1 - bin_img, ksize)
return out
if __name__ == '__main__':
I = cv2.imread('bev_img.jpg', 0)
I_torch = torch.from_numpy(I // 255).type(torch.float32).unsqueeze(0).unsqueeze(0)
I_dilate = dilate(I_torch)
I_dilate = (I_dilate.detach().squeeze().cpu().numpy() * 255).astype('uint8')
I_erode = erode(I_torch, ksize=10)
I_erode = (I_erode.detach().squeeze().cpu().numpy()*255).astype('uint8')
show = np.concatenate((I, I_erode, I_dilate), axis=1)
# cv2.imshow('show', show)
# cv2.waitKey(0)
cv2.imwrite('torch_erode_bev_img.jpg', show)
效果图
- 依次是:原图,膨胀,腐蚀