File size: 1,962 Bytes
d85fde4 1eb87a5 4a3fe77 d85fde4 4a3fe77 d85fde4 4a3fe77 d85fde4 4a3fe77 |
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 |
from functools import partial
import jax
import numpy as np
def repeat_vmap(fun, in_axes=None):
if in_axes is None:
in_axes = [0]
for axes in in_axes:
fun = jax.vmap(fun, in_axes=axes)
return fun
def make_grid(patch_size: int | tuple[int, int]):
if isinstance(patch_size, int):
patch_size = (patch_size, patch_size)
offset_h, offset_w = 1 / (2 * np.array(patch_size))
space_h = np.linspace(-0.5 + offset_h, 0.5 - offset_h, patch_size[0])
space_w = np.linspace(-0.5 + offset_w, 0.5 - offset_w, patch_size[1])
return np.stack(np.meshgrid(space_h, space_w, indexing='ij'), axis=-1) # [h, w]
def interpolate_grid(coords, grid, order=0):
"""
Args:
coords: Tensor de shape (B, H, W, 2) ou (H, W, 2)
grid: Tensor de shape (B, H', W', C)
"""
# Adicionar dimensão de batch se necessário
if coords.ndim == 3:
coords = coords[np.newaxis, ...]
# Verificar dimensões
assert coords.ndim == 4, f"Dimensões inválidas para coords: {coords.shape}"
assert grid.ndim == 4, f"Dimensões inválidas para grid: {grid.shape}"
# Ajustar transposição de forma segura
try:
coords = coords.transpose((0, 3, 1, 2))
except ValueError as e:
raise ValueError(f"Falha na transposição: {coords.shape} → (0,3,1,2)") from e
# Conversão de coordenadas
coords = coords.at[:, 0].set(coords[:, 0] * grid.shape[-3] + (grid.shape[-3] - 1) / 2)
coords = coords.at[:, 1].set(coords[:, 1] * grid.shape[-2] + (grid.shape[-2] - 1) / 2)
# Interpolação com JAX
map_coordinates = partial(jax.scipy.ndimage.map_coordinates,
order=order,
mode='nearest')
return jax.vmap( # Sobre batches
jax.vmap( # Sobre canais
map_coordinates,
in_axes=(2, None), # (C, H', W'), (B, 2, H, W)
out_axes=2
)
)(grid, coords) |