CondViT-LRVSF-Demo / src /transform.py
Slep's picture
Initial Demo.
c45703e
raw
history blame
842 Bytes
from torchvision.transforms import transforms as tf
import torchvision.transforms.functional as F
class SquarePad:
def __init__(self, color):
self.col = color
def __call__(self, image):
max_wh = max(image.size)
p_left, p_top = [(max_wh - s) // 2 for s in image.size]
p_right, p_bottom = [
max_wh - (s + pad) for s, pad in zip(image.size, [p_left, p_top])
]
padding = (p_left, p_top, p_right, p_bottom)
return F.pad(image, padding, self.col, "constant")
def valid_tf(size):
return tf.Compose(
[
SquarePad(255),
tf.Resize(size),
tf.ToTensor(),
tf.Normalize(
mean=(0.48145466, 0.4578275, 0.40821073),
std=(0.26862954, 0.26130258, 0.27577711),
),
]
)