boxfilter for pytorch

boxfilter.py

def _diff_x(src: Tensor, w: int) -> Tensor:
    cum_src = src.cumsum(-2)

    left = cum_src[..., w : 2 * w + 1, :]
    middle = cum_src[..., 2 * w + 1 :, :] - cum_src[..., : -2 * w - 1, :]
    right = cum_src[..., -1:, :] - cum_src[..., -2 * w - 1 : -w - 1, :]

    return torch.vstack((left, middle, right))


def _diff_y(src: Tensor, w: int) -> Tensor:
    cum_src = src.cumsum(-1)

    left = cum_src[..., w : 2 * w + 1]
    middle = cum_src[..., 2 * w + 1 :] - cum_src[..., : -2 * w - 1]
    right = cum_src[..., -1:] - cum_src[..., -2 * w - 1 : -w - 1]

    return torch.hstack((left, middle, right))


def _box_filter2d(src: Tensor, wing: int) -> Tensor:
    return _diff_y(_diff_x(src, wing), wing)