Spaces:
Running
Running
File size: 2,010 Bytes
80ebcb3 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from typing import List, Literal, Tuple
import torch
import torch.nn.functional as F
def center_crop_image(image: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
num_channels, height, width = image.shape
crop_h, crop_w = size
top = (height - crop_h) // 2
left = (width - crop_w) // 2
return image[:, top : top + crop_h, left : left + crop_w]
def resize_crop_image(image: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
num_channels, height, width = image.shape
target_h, target_w = size
scale = max(target_h / height, target_w / width)
new_h, new_w = int(height * scale), int(width * scale)
image = F.interpolate(image, size=(new_h, new_w), mode="bilinear", align_corners=False)
return center_crop_image(image, size)
def bicubic_resize_image(image: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
return F.interpolate(image, size=size, mode="bicubic", align_corners=False)
def find_nearest_resolution_image(image: torch.Tensor, resolution_buckets: List[Tuple[int, int]]) -> Tuple[int, int]:
num_channels, height, width = image.shape
aspect_ratio = width / height
def aspect_ratio_diff(bucket):
return abs((bucket[1] / bucket[0]) - aspect_ratio)
return min(resolution_buckets, key=aspect_ratio_diff)
def resize_to_nearest_bucket_image(
image: torch.Tensor,
resolution_buckets: List[Tuple[int, int]],
resize_mode: Literal["center_crop", "resize_crop", "bicubic"] = "bicubic",
) -> torch.Tensor:
target_size = find_nearest_resolution_image(image, resolution_buckets)
if resize_mode == "center_crop":
return center_crop_image(image, target_size)
elif resize_mode == "resize_crop":
return resize_crop_image(image, target_size)
elif resize_mode == "bicubic":
return bicubic_resize_image(image, target_size)
else:
raise ValueError(
f"Invalid resize_mode: {resize_mode}. Choose from 'center_crop', 'resize_crop', or 'bicubic'."
)
|