File size: 15,953 Bytes
c69c4af |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
import math
from functools import reduce
from logging import getLogger
from typing import Any, Callable, Tuple
import numpy as np
import torch
import torchvision.transforms as tv
from PIL import Image
from torchvision.transforms import functional as F
from torchvision.transforms.functional import InterpolationMode
logger = getLogger()
MEAN = (0.5, 0.5, 0.5)
STD = (0.5, 0.5, 0.5)
def get_image_transform(
vision_input_type: str = "vanilla",
image_res: int = 336,
max_num_tiles: int = 1,
normalize_img: bool = True,
) -> Tuple[Callable, int]:
if vision_input_type == "thumb+tile":
transforms = VariableSizeImageTransform(
size=image_res,
max_num_tiles=max_num_tiles,
normalize_img=normalize_img,
use_thumbnail="before",
)
else:
transforms = ImageTransform(
size=image_res,
normalize_img=normalize_img,
)
logger.info(
f"Initalized transforms with: vision_input_type: '{vision_input_type}' and max_num_tiles: {max_num_tiles}."
)
return transforms
class ImageTransform(object):
"""
Image transform will resize the longer edge to a given size and pad the shorter edge with mean pixel value of the image.
"""
def __init__(
self,
size: int = 336,
normalize_img: bool = True,
) -> None:
self.size = size
self._mean = MEAN
self._std = STD
logger.info(f"ImageTransform size: {self.size}")
self.to_tensor = tv.ToTensor()
self.normalize = (
tv.Normalize(
mean=self._mean,
std=self._std,
inplace=True,
)
if normalize_img
else lambda x: x
)
def __call__(self, image: Image.Image):
w, h = image.size
image = F.resize(
image, (self.size, self.size), interpolation=InterpolationMode.BICUBIC
)
image = self.to_tensor(image)
image = self.normalize(image)
# Add chunk dim to make it compatible with existing dataloaders
image = image.view(1, 3, self.size, self.size)
return image, (w, h)
def _transform_torch_tensor(self, image: torch.Tensor):
h, w = image.shape[-2:] # Image shape (C, H, W) or (N, C, H, W)
image = F.resize(
image, size=(self.size, self.size), interpolation=InterpolationMode.BICUBIC
)
image = (
image.to(torch.float32) / 255.0
) # Convert to float and scale to [0, 1] range
image = self.normalize(image)
return image, (w, h)
class VariableSizeImageTransform(object):
"""
The variable size image transform will resize the image dynamically
based on the image aspect ratio and the number of image chunks we allow.
The algorithm will not upsample low-res images to fit a certain aspect
ratio, because that leads to a significant degradation in image quality.
For example, if an input image is of size 300x800, and we want to allow
a maximum of 16 image chunks, it will find the closest aspect ratio that
is allowed within 16 image chunks, i.e., 2:5 = 2 horizontal patches and
5 vertical patches, giving a total of 10 chunks.
The image will then be resized to products of the base size (default is
224px because MetaCLIP takes that), so in this case it will be resized to
2*224:5*224 = 448:1120, where we maintain the original aspect ratio and
pad with the mean value for the rest. This approach minimizes the amount
of padding required for any arbitrary resolution.
The final output will therefore be of shape (11, 3, 224, 224), where 10
patches are coming from the resizing and chunking, and the first patch
is a downsampled version of the image that preserves aspect ratios.
"""
def __init__(
self,
size: int = 336,
normalize_img: bool = True,
max_num_tiles: int = 1,
use_thumbnail: str = "no",
area_limit: bool = False,
) -> None:
self.size = size
self._mean = MEAN
self._std = STD
logger.info(f"VariableSizeImageTransform size: {self.size}")
self.to_tensor = tv.ToTensor()
self.normalize = (
tv.Normalize(
mean=self._mean,
std=self._std,
inplace=True,
)
if normalize_img
else lambda x: x
)
self.area_limit = area_limit
self.max_num_tiles = max_num_tiles
self.use_thumbnail = use_thumbnail
if self.use_thumbnail != "no":
self.thumbnail_transform = ImageTransform(
size=self.size,
normalize_img=normalize_img,
)
@staticmethod
def _factors(n: int):
"""Return all factors of a number."""
return set(
reduce(
list.__add__,
([i, n // i] for i in range(1, int(n**0.5) + 1) if n % i == 0),
)
)
def _find_supported_aspect_ratios(self):
"""
This function computes all the allowed aspect ratios for a fixed
number of input chunks.
For example, with `num_tiles=5`, it will return:
{
0.2: [(1, 5)],
5.0: [(5, 1)],
0.25: [(1, 4)],
1.0: [(2, 2), (1, 1)],
4.0: [(4, 1)],
0.3333333333333333: [(1, 3)],
3.0: [(3, 1)],
0.5: [(1, 2)],
2.0: [(2, 1)]
}
"""
asp_dict = {}
for chunk_size in range(self.max_num_tiles, 0, -1):
_factors = sorted(VariableSizeImageTransform._factors(chunk_size))
_asp_ratios = [(x, chunk_size // x) for x in _factors]
for ratio in _asp_ratios:
k = ratio[0] / ratio[1]
if k not in asp_dict:
asp_dict[k] = [ratio]
else:
asp_dict[k].append(ratio)
return asp_dict
def _find_closest_aspect_ratio(self, img_width: int, img_height: int) -> Tuple:
"""
Given an image width, height and target number of chunks
this function will find the closest supported aspect ratio.
"""
tgt_ar = img_width / img_height
asp_dict = self._find_supported_aspect_ratios()
cl_d, cl_p = 1e23, None
if tgt_ar >= 1:
cl_p = min(
[k for k in asp_dict.keys() if k <= tgt_ar],
key=lambda x: abs(x - tgt_ar),
)
v = asp_dict[cl_p]
# select width
widths = [(idx, self.size * vv[0]) for idx, vv in enumerate(v)]
tgt_idx = max(widths, key=lambda x: x[1])[0]
else:
cl_p = min(
[k for k in asp_dict.keys() if k > tgt_ar],
key=lambda x: abs(1 / x - 1 / tgt_ar),
)
v = asp_dict[cl_p]
# select height
heights = [(idx, self.size * vv[1]) for idx, vv in enumerate(v)]
tgt_idx = max(heights, key=lambda x: x[1])[0]
out = v[tgt_idx]
return out
def _resize(
self, image: Image.Image, target_width: int, target_height: int
) -> Image.Image:
# Resize longer edge to given size.
w, h = image.size
scale = w / h
if scale > 1.0:
# width > height
new_w = target_width
new_h = math.floor(new_w / scale)
else:
# height >= width
new_h = target_height
new_w = math.floor(new_h * scale)
image = F.resize(image, (new_h, new_w))
return image
def _pad(self, image: Image.Image, new_width: int, new_height: int) -> Image.Image:
mean_per_channel = tuple(
np.clip(np.array(image).mean(axis=(0, 1)), 0, 255).astype(np.uint8)
)
new_im = Image.new(mode="RGB", size=(new_width, new_height), color=(0, 0, 0)) # type: ignore
new_im.paste(image)
return new_im
def _split(self, image: torch.Tensor, ncw: int, nch: int) -> torch.Tensor:
# Split image into number of required tiles (width x height)
num_channels, height, width = image.size()
image = image.view(num_channels, nch, height // nch, ncw, width // ncw)
# Permute dimensions to reorder the axes
image = image.permute(1, 3, 0, 2, 4).contiguous()
# Reshape into the desired output shape (batch_size * 4, num_channels, width/2, height/2)
image = image.view(ncw * nch, num_channels, height // nch, width // ncw)
return image
def _get_image_height_width(
self, image_width: int, image_height: int, target_width: int, target_height: int
) -> Tuple[int, int]:
"""
Given image width, height and target width, height for the canvas, return the dimensions of how the image would be resized
with aspect ratio preservation.
"""
scale = image_width / image_height
if scale > 1.0:
# Width is larger than height
# Rescaling factor is the minimum of the two scaling factors. Else one side would be outside of the canvas.
rescaling_factor = min(
target_width / image_width, target_height / image_height
)
# Set new width to target width and height to the rescaled height.
new_w = rescaling_factor * image_width
new_h = math.floor(new_w / scale)
else:
# Height is larger than width
# Rescaling factor is the minimum of the two scaling factors. Else one side would be outside of the canvas.
rescaling_factor = min(
target_width / image_width, target_height / image_height
)
# Set new height to target height and width to the rescaled width.
new_h = rescaling_factor * image_height
new_w = math.floor(new_h * scale)
return new_w, new_h
def _fit_image_to_canvas(
self, img_width: int, img_height: int, area_limit: bool
) -> Any:
"""
Given an image width, height and target number of chunks this function will see if the image
can be fit into any of the canvases that can be build from arranging the tiles in a grid.
If the image can be fit onto several canvases, it will return the canvas where the shorter edge
of the image will be largest.
If area_limit is set to True, the tie-breaking prefers the canvas where area is less than 2x the original area.
"""
# Initialize the optimal canvas to None. If no canvas is found where image fits, function returns None.
optimal_canvas = None
optimal_image_width_height = None
scale = img_width / img_height
# Gather all potential supported image resolutions and iterate through them to find best match
potential_arrangements = [
item
for sublist in self._find_supported_aspect_ratios().values()
for item in sublist
]
for n_w, n_h in potential_arrangements:
# Compute the canvas size
canvas_width, canvas_height = n_w * self.size, n_h * self.size
# Check if image can fit into the canvas without downsampling
if canvas_width >= img_width and canvas_height >= img_height:
# If we did not find a good canvas yet, we will use the current one
if optimal_canvas is None:
# Set optimal canvas and determine the actual image height and width in the canvas with aspect ratio preserving resampling
optimal_canvas = (n_w, n_h)
optimal_image_width_height = self._get_image_height_width(
image_width=img_width,
image_height=img_height,
target_width=n_w * self.size,
target_height=n_h * self.size,
)
else:
# If we already found an optimal canvas before, we will check if the shorter edge of the image will be larger than the current optimal canvas.
# This means we can potentially upsample the image resolution which is beneficial to performance.
image_width_height = self._get_image_height_width(
image_width=img_width,
image_height=img_height,
target_width=n_w * self.size,
target_height=n_h * self.size,
)
if area_limit:
# Prioritize aspect ratio, and choose best within area limit when tied.
curr_scale = image_width_height[0] / image_width_height[1]
optim_scale = (
optimal_image_width_height[0]
/ optimal_image_width_height[1]
)
if abs(scale - curr_scale) < abs(scale - optim_scale):
# 1. optimize aspect ratio
optimal_canvas = (n_w, n_h)
optimal_image_width_height = image_width_height
elif abs(scale - curr_scale) == abs(scale - optim_scale):
# 2. optimize area
if (
image_width_height[0] * image_width_height[1]
< 2 * img_width * img_height
):
# 2.1 area is less than 2x the original area
optimal_canvas = (n_w, n_h)
optimal_image_width_height = image_width_height
else:
# NOTE: L3V dynamid tiling. Priortize biggest canvas.
if (
scale < 1.0
and (image_width_height[0] >= optimal_image_width_height[0])
) or (
scale >= 1.0
and (image_width_height[1] >= optimal_image_width_height[1])
):
optimal_canvas = (n_w, n_h)
optimal_image_width_height = image_width_height
return optimal_canvas
def __call__(self, image: Image.Image) -> Tuple[Any, Any]:
assert isinstance(image, Image.Image), type(image)
if self.use_thumbnail != "no":
thumbnail = self.thumbnail_transform(image)[0]
w, h = image.size
# Check if the image can be fit to the canvas without downsampling
ar = self._fit_image_to_canvas(
img_width=w, img_height=h, area_limit=self.area_limit
)
if ar is None:
# If we did not find a canvas, we have to find the closest aspect ratio and downsample the image
ar = self._find_closest_aspect_ratio(img_width=w, img_height=h)
image = F.resize(
image,
(ar[1] * self.size, ar[0] * self.size), # (h, w)
interpolation=InterpolationMode.BICUBIC,
)
image = self._pad(image, ar[0] * self.size, ar[1] * self.size)
image = self.to_tensor(image)
image = self.normalize(image)
image = self._split(image, ar[0], ar[1]) # type: ignore
if self.use_thumbnail == "before":
image = torch.cat((thumbnail, image), dim=0)
elif self.use_thumbnail == "after":
image = torch.cat((image, thumbnail), dim=0)
elif self.use_thumbnail == "both":
image = torch.cat((thumbnail, image, thumbnail), dim=0)
return image, ar
|