| import torch |
| import numpy as np |
| import torchvision |
| from PIL import Image |
| from torchvision.transforms.functional import InterpolationMode |
| import torchvision.transforms as transforms |
|
|
| def padding_336(b): |
| width, height = b.size |
| tar = int(np.ceil(height / 336) * 336) |
| top_padding = int((tar - height)/2) |
| bottom_padding = tar - height - top_padding |
| left_padding = 0 |
| right_padding = 0 |
| b = transforms.functional.pad(b, [left_padding, top_padding, right_padding, bottom_padding], fill=[255,255,255]) |
|
|
| return b |
|
|
| def HD_transform(img, hd_num=16): |
| width, height = img.size |
| trans = False |
| if width < height: |
| img = img.transpose(Image.TRANSPOSE) |
| trans = True |
| width, height = img.size |
| ratio = (width/ height) |
| scale = 1 |
| while scale*np.ceil(scale/ratio) <= hd_num: |
| scale += 1 |
| scale -= 1 |
| new_w = int(scale * 336) |
| new_h = int(new_w / ratio) |
|
|
| img = transforms.functional.resize(img, [new_h, new_w],) |
| img = padding_336(img) |
| width, height = img.size |
| if trans: |
| img = img.transpose(Image.TRANSPOSE) |
|
|
| return img |
|
|
|
|
|
|
