Pytorch 张量滤波中值替换

使用一个小尺寸的值都为1张量对一个大尺寸的张量进行中值滤波，将小尺寸张量对应卷乘大尺寸张量，并将卷乘得到的结果排序后将中值替换：

import torch
import numpy as np

def medfilt(volume):
    kernel_size = torch.ones(3, 3)
    w, h = volume.shape[0], volume.shape[1]
    for i in range(h- 3 + 1):
        for j in range(w - 3 + 1):
            interim = kernel_size * volume[i:i+3, j:j+3]
            interim, indices = torch.sort(interim)
            print('3*3:',interim)
            volume[i+1][j+1] = interim[1][1]

    return volume
    
if __name__ == '__main__':
    # 创建一个形如(B, C, H, W)的张量，对应pytorch神经网络图片输入
    a = torch.randint(2, 9, (2, 4, 5, 5))

    # 按 B 读取，即一张张读取图片
    for i in range(a.shape[0]):
        # a[i, 0, ...]相当于读入的第i张图像，第0维的图像
        print('input:', a[i, 0, ...])
        a[i, 0, ...] = medfilt(a[i, 0, ...])
        print('output:', a[i, 0, ...])