| try: |
| import jitfields |
| available = True |
| except (ImportError, ModuleNotFoundError): |
| jitfields = None |
| available = False |
| from .utils import make_list |
| import torch |
|
|
|
|
| def first2last(input, ndim): |
| insert = input.dim() <= ndim |
| if insert: |
| input = input.unsqueeze(-1) |
| else: |
| input = torch.movedim(input, -ndim-1, -1) |
| return input, insert |
|
|
|
|
| def last2first(input, ndim, inserted, grad=False): |
| if inserted: |
| input = input.squeeze(-1 - grad) |
| else: |
| input = torch.movedim(input, -1 - grad, -ndim-1 - grad) |
| return input |
|
|
|
|
| def grid_pull(input, grid, interpolation='linear', bound='zero', |
| extrapolate=False, prefilter=False): |
| ndim = grid.shape[-1] |
| input, inserted = first2last(input, ndim) |
| input = jitfields.pull(input, grid, order=interpolation, bound=bound, |
| extrapolate=extrapolate, prefilter=prefilter) |
| input = last2first(input, ndim, inserted) |
| return input |
|
|
|
|
| def grid_push(input, grid, shape=None, interpolation='linear', bound='zero', |
| extrapolate=False, prefilter=False): |
| ndim = grid.shape[-1] |
| input, inserted = first2last(input, ndim) |
| input = jitfields.push(input, grid, shape, order=interpolation, bound=bound, |
| extrapolate=extrapolate, prefilter=prefilter) |
| input = last2first(input, ndim, inserted) |
| return input |
|
|
|
|
| def grid_count(grid, shape=None, interpolation='linear', bound='zero', |
| extrapolate=False): |
| return jitfields.count(grid, shape, order=interpolation, bound=bound, |
| extrapolate=extrapolate) |
|
|
|
|
| def grid_grad(input, grid, interpolation='linear', bound='zero', |
| extrapolate=False, prefilter=False): |
| ndim = grid.shape[-1] |
| input, inserted = first2last(input, ndim) |
| input = jitfields.grad(input, grid, order=interpolation, bound=bound, |
| extrapolate=extrapolate, prefilter=prefilter) |
| input = last2first(input, ndim, inserted, True) |
| return input |
|
|
|
|
| def spline_coeff(input, interpolation='linear', bound='dct2', dim=-1, |
| inplace=False): |
| func = jitfields.spline_coeff_ if inplace else jitfields.spline_coeff |
| return func(input, interpolation, bound=bound, dim=dim) |
|
|
|
|
| def spline_coeff_nd(input, interpolation='linear', bound='dct2', dim=None, |
| inplace=False): |
| func = jitfields.spline_coeff_nd_ if inplace else jitfields.spline_coeff_nd |
| return func(input, interpolation, bound=bound, ndim=dim) |
|
|
|
|
| def resize(image, factor=None, shape=None, anchor='c', |
| interpolation=1, prefilter=True, **kwargs): |
| kwargs.setdefault('bound', 'nearest') |
| ndim = max(len(make_list(factor or [])), |
| len(make_list(shape or [])), |
| len(make_list(anchor or []))) or (image.dim() - 2) |
| return jitfields.resize(image, factor=factor, shape=shape, ndim=ndim, |
| anchor=anchor, order=interpolation, |
| bound=kwargs['bound'], prefilter=prefilter) |
|
|
|
|
| def restrict(image, factor=None, shape=None, anchor='c', |
| interpolation=1, reduce_sum=False, **kwargs): |
| kwargs.setdefault('bound', 'nearest') |
| ndim = max(len(make_list(factor or [])), |
| len(make_list(shape or [])), |
| len(make_list(anchor or []))) or (image.dim() - 2) |
| return jitfields.restrict(image, factor=factor, shape=shape, ndim=ndim, |
| anchor=anchor, order=interpolation, |
| bound=kwargs['bound'], reduce_sum=reduce_sum) |
|
|
