pytorch-image-models

1
""" Median Pool
2
Hacked together by / Copyright 2020 Ross Wightman
3
"""
4
import torch.nn as nn
5
import torch.nn.functional as F
6
from .helpers import to_2tuple, to_4tuple
7

8

9
class MedianPool2d(nn.Module):
10
    """ Median pool (usable as median filter when stride=1) module.
11

12
    Args:
13
         kernel_size: size of pooling kernel, int or 2-tuple
14
         stride: pool stride, int or 2-tuple
15
         padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad
16
         same: override padding and enforce same padding, boolean
17
    """
18
    def __init__(self, kernel_size=3, stride=1, padding=0, same=False):
19
        super(MedianPool2d, self).__init__()
20
        self.k = to_2tuple(kernel_size)
21
        self.stride = to_2tuple(stride)
22
        self.padding = to_4tuple(padding)  # convert to l, r, t, b
23
        self.same = same
24

25
    def _padding(self, x):
26
        if self.same:
27
            ih, iw = x.size()[2:]
28
            if ih % self.stride[0] == 0:
29
                ph = max(self.k[0] - self.stride[0], 0)
30
            else:
31
                ph = max(self.k[0] - (ih % self.stride[0]), 0)
32
            if iw % self.stride[1] == 0:
33
                pw = max(self.k[1] - self.stride[1], 0)
34
            else:
35
                pw = max(self.k[1] - (iw % self.stride[1]), 0)
36
            pl = pw // 2
37
            pr = pw - pl
38
            pt = ph // 2
39
            pb = ph - pt
40
            padding = (pl, pr, pt, pb)
41
        else:
42
            padding = self.padding
43
        return padding
44

45
    def forward(self, x):
46
        x = F.pad(x, self._padding(x), mode='reflect')
47
        x = x.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])
48
        x = x.contiguous().view(x.size()[:4] + (-1,)).median(dim=-1)[0]
49
        return x
50