| import torch.nn.functional as F | |
| class InputPadder: | |
| """ Pads images such that dimensions are divisible by divisor """ | |
| def __init__(self, dims, divisor=16): | |
| self.ht, self.wd = dims[-2:] | |
| pad_ht = (((self.ht // divisor) + 1) * divisor - self.ht) % divisor | |
| pad_wd = (((self.wd // divisor) + 1) * divisor - self.wd) % divisor | |
| self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, pad_ht // 2, pad_ht - pad_ht // 2] | |
| def pad(self, *inputs): | |
| if len(inputs) == 1: | |
| return F.pad(inputs[0], self._pad, mode='replicate') | |
| else: | |
| return [F.pad(x, self._pad, mode='replicate') for x in inputs] | |
| def unpad(self, *inputs): | |
| if len(inputs) == 1: | |
| return self._unpad(inputs[0]) | |
| else: | |
| return [self._unpad(x) for x in inputs] | |
| def _unpad(self, x): | |
| ht, wd = x.shape[-2:] | |
| c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] | |
| return x[..., c[0]:c[1], c[2]:c[3]] | |