25 lines
785 B
Python
25 lines
785 B
Python
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):
|
|
return [F.pad(x, self._pad, mode="replicate") 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]]
|