69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
""" Interpolation helpers for timm layers
|
|
|
|
RegularGridInterpolator from https://github.com/sbarratt/torch_interpolations
|
|
Copyright Shane Barratt, Apache 2.0 license
|
|
"""
|
|
import torch
|
|
from itertools import product
|
|
|
|
|
|
class RegularGridInterpolator:
|
|
""" Interpolate data defined on a rectilinear grid with even or uneven spacing.
|
|
Produces similar results to scipy RegularGridInterpolator or interp2d
|
|
in 'linear' mode.
|
|
|
|
Taken from https://github.com/sbarratt/torch_interpolations
|
|
"""
|
|
|
|
def __init__(self, points, values):
|
|
self.points = points
|
|
self.values = values
|
|
|
|
assert isinstance(self.points, tuple) or isinstance(self.points, list)
|
|
assert isinstance(self.values, torch.Tensor)
|
|
|
|
self.ms = list(self.values.shape)
|
|
self.n = len(self.points)
|
|
|
|
assert len(self.ms) == self.n
|
|
|
|
for i, p in enumerate(self.points):
|
|
assert isinstance(p, torch.Tensor)
|
|
assert p.shape[0] == self.values.shape[i]
|
|
|
|
def __call__(self, points_to_interp):
|
|
assert self.points is not None
|
|
assert self.values is not None
|
|
|
|
assert len(points_to_interp) == len(self.points)
|
|
K = points_to_interp[0].shape[0]
|
|
for x in points_to_interp:
|
|
assert x.shape[0] == K
|
|
|
|
idxs = []
|
|
dists = []
|
|
overalls = []
|
|
for p, x in zip(self.points, points_to_interp):
|
|
idx_right = torch.bucketize(x, p)
|
|
idx_right[idx_right >= p.shape[0]] = p.shape[0] - 1
|
|
idx_left = (idx_right - 1).clamp(0, p.shape[0] - 1)
|
|
dist_left = x - p[idx_left]
|
|
dist_right = p[idx_right] - x
|
|
dist_left[dist_left < 0] = 0.
|
|
dist_right[dist_right < 0] = 0.
|
|
both_zero = (dist_left == 0) & (dist_right == 0)
|
|
dist_left[both_zero] = dist_right[both_zero] = 1.
|
|
|
|
idxs.append((idx_left, idx_right))
|
|
dists.append((dist_left, dist_right))
|
|
overalls.append(dist_left + dist_right)
|
|
|
|
numerator = 0.
|
|
for indexer in product([0, 1], repeat=self.n):
|
|
as_s = [idx[onoff] for onoff, idx in zip(indexer, idxs)]
|
|
bs_s = [dist[1 - onoff] for onoff, dist in zip(indexer, dists)]
|
|
numerator += self.values[as_s] * \
|
|
torch.prod(torch.stack(bs_s), dim=0)
|
|
denominator = torch.prod(torch.stack(overalls), dim=0)
|
|
return numerator / denominator
|