Pytorch——基于max pooling实现图像腐蚀和膨胀

发现一个宝藏回答：

• 用 Pytorch 实现图像的腐蚀

代码

• 可以通过调节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)

效果图

• 依次是：原图，膨胀，腐蚀