Computing aspect ratio buckets

arb.py

import numpy as np
import math
from numpy.typing import NDArray

# we are trying to make buckets of varying aspect ratios,
# all with about the same area (equivalent to a 512x512 square)
square_side = 512
buckets = 8

widest_aspect: float = math.atan2(1, 2)  # 1/2 = 0.5 aspect ratio
tallest_aspect: float = math.atan2(2, 1) # 2/1 = 2.0 aspect ratio

thetas: NDArray = np.linspace(widest_aspect, tallest_aspect, buckets)

hypotenuses: NDArray = square_side/(2*np.sin(thetas)*np.cos(thetas))**.5

heights: NDArray = hypotenuses*2**.5*np.sin(thetas)
widths: NDArray = hypotenuses*2**.5*np.cos(thetas)

buckets: NDArray = np.column_stack([heights, widths])

# all buckets will have about the same area as square_side**2
# you can verify this by multiplying their widths by their heights:
# np.multiply.reduce(buckets, axis=-1)

# round to nearest 8px (e.g. to be suitable for encoding in VAE)
buckets_rounded: NDArray = 8*np.rint(buckets / 8).astype(int)
pass