| 12345678910111213141516171819202122232425262728 |
- import torch.nn.functional as F
- class InputPadder:
- """Pads images such that dimensions are divisible by 8"""
- # TODO: Ideally, this should be part of the eval transforms preset, instead
- # of being part of the validation code. It's not obvious what a good
- # solution would be, because we need to unpad the predicted flows according
- # to the input images' size, and in some datasets (Kitti) images can have
- # variable sizes.
- def __init__(self, dims, mode="sintel"):
- self.ht, self.wd = dims[-2:]
- pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8
- pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8
- if mode == "sintel":
- self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, pad_ht // 2, pad_ht - pad_ht // 2]
- else:
- self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht]
- 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]]
|
