torch-fourier-shift 0.1.0

Shift 2D/3D images by phase shifting Fourier transforms in PyTorch

torch-phase-shift

torch-fourier-shift is a package for shifting 1D, 2D and 3D images with subpixel precision by applying phase shifts to Fourier transforms in PyTorch.

import torch
from torch_fourier_shift import fourier_shift_image_2d

# create a dummy image
my_image = torch.tensor(
    [[0, 0, 0, 0, 0, 0],
     [0, 1, 0, 0, 0, 0],
     [0, 0, 0, 0, 0, 0],
     [0, 0, 0, 0, 0, 0],
     [0, 0, 0, 0, 0, 0],
     [0, 0, 0, 0, 0, 0]]
)

# shift the image by 1 pixel in dim 0, 2 pixels in dim 1
shifted_image = fourier_shift_image_2d(image=my_image, shifts=torch.tensor([1, 2]))

API's are equivalent for 1D and 3D images.

Installation

torch-fourier-shift is available on PyPI.

pip install torch-fourier-shift

Usage

Please check the the docs at teamtomo.org/torch-fourier-shift

Caching

Some functions are equipped with an argument called cache_intermediates. If you set cache_intermediates=True, an LRU cache will be used to avoid recomputing intermediate results. Note that this might affect gradient calculations.

By default, the size of the cache is 3, and can be changed with an environmental variable called TORCH_FOURIER_SHIFT_CACHE_SIZE. Just do

export TORCH_FOURIER_SHIFT_CACHE_SIZE=5

or

os.environ["TORCH_FOURIER_SHIFT_CACHE_SIZE"]=5

before importing the torch_fourier_shift module.