File size: 725 Bytes
7896def | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | import torch
import torch.nn as nn
import torchvision.transforms as TF
class ToTensor(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor):
assert x.dtype == torch.uint8
x = x.to(torch.float32) / 255.0
return x
class Pad(nn.Module):
def __init__(self, padding, fill=0, padding_mode='constant'):
super().__init__()
self.padding = padding
self.fill = fill
self.padding_mode = padding_mode
self.pad = TF.Pad(padding=tuple(padding), fill=fill, padding_mode=padding_mode)
def forward(self, x: torch.Tensor):
assert x.ndim == 4, "Can only pad tensor of 4 dims."
return self.pad(x)
|