File size: 1,895 Bytes
d85fde4 1eb87a5 d160dc6 1eb87a5 054a11a d85fde4 054a11a d85fde4 054a11a d160dc6 054a11a d160dc6 054a11a d85fde4 054a11a d160dc6 054a11a d160dc6 054a11a 4a3fe77 d160dc6 4a3fe77 054a11a d160dc6 4a3fe77 d160dc6 054a11a |
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 |
from functools import partial
import jax
import jax.numpy as jnp
import numpy as np
def repeat_vmap(fun, 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]):
"""Gera grid de coordenadas com segurança numérica"""
# Garantir tamanho mínimo de 8x8
if isinstance(patch_size, int):
h = w = max(8, patch_size)
else:
h, w = (max(8, ps) for ps in patch_size)
# Espaçamento preciso entre pontos
y_space = np.linspace(-0.5 + 1 / (2 * h), 0.5 - 1 / (2 * h), h)
x_space = np.linspace(-0.5 + 1 / (2 * w), 0.5 - 1 / (2 * w), w)
# Criar grid com dimensões (1, H, W, 2)
grid = np.stack(np.meshgrid(y_space, x_space, indexing='ij'), axis=-1)
return grid[np.newaxis, ...]
def interpolate_grid(coords, grid, order=0):
"""Interpolação segura com verificação de dimensões"""
try:
# Converter para JAX array e validar formato
coords = jnp.asarray(coords)
if coords.ndim != 4 or coords.shape[-1] != 2:
raise ValueError(
f"Dimensões inválidas: {coords.shape}. Esperado (B, H, W, 2)"
)
# Transformação de coordenadas
coords = coords.transpose((0, 3, 1, 2))
coords = coords.at[:, 0].set(
coords[:, 0] * (grid.shape[-3] - 1) + (grid.shape[-3] - 1) / 2
)
coords = coords.at[:, 1].set(
coords[:, 1] * (grid.shape[-2] - 1) + (grid.shape[-2] - 1) / 2
)
# Interpolação vetorizada
map_fn = jax.vmap(jax.vmap(
partial(jax.scipy.ndimage.map_coordinates, order=order, mode='nearest'),
in_axes=(2, None),
out_axes=2
))
return map_fn(grid, coords)
except Exception as e:
raise RuntimeError(f"Erro de interpolação: {str(e)}") from e |