31 lines
888 B
Python
31 lines
888 B
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def preprocess_image(image: torch.Tensor, device, factor=32, min_size=225):
|
|
# [B, H, W, C]
|
|
input_img = image.permute(0, 3, 1, 2).to(device) # [B, C, H, W]
|
|
_, _, h, w = input_img.shape
|
|
|
|
# Resize if either dimension is smaller than min_size
|
|
if h < min_size or w < min_size:
|
|
scale = min_size / min(h, w)
|
|
new_h, new_w = int(h * scale), int(w * scale)
|
|
input_img = F.interpolate(
|
|
input_img, size=(new_h, new_w), mode="bilinear", align_corners=False
|
|
)
|
|
h, w = new_h, new_w
|
|
|
|
# Calculate padding
|
|
H = ((h + factor - 1) // factor) * factor
|
|
W = ((w + factor - 1) // factor) * factor
|
|
|
|
padh = H - h
|
|
padw = W - w
|
|
|
|
# Pad the image
|
|
if padh > 0 or padw > 0:
|
|
input_img = F.pad(input_img, (0, padw, 0, padh), mode="reflect")
|
|
|
|
return input_img, h, w
|