Coverage for src/susi/utils/sub_shift.py: 95%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4Util for sub pixel shifting 

5 

6@author: hoelken 

7""" 

8import numpy as np 

9from scipy import ndimage 

10 

11 

12class Shifter: 

13 """ 

14 Helper class to shift images in subpixel amounts using FFT 

15 """ 

16 

17 @staticmethod 

18 def d1shift(img, offset): 

19 """ 

20 Creates a FFTd version of the image and shifts it by the offset amount. 

21 

22 WARNING The shift applied is -offset, thus a positive value of offset means a shift torwards 0 

23 

24 ## Params 

25 - `ìmg`: An 1d array object representing the image row 

26 - `offset`: The amount the center is shifted must (be float) 

27 """ 

28 ft_img = np.fft.fft(img) 

29 result = ndimage.fourier_shift(ft_img, shift=-offset) 

30 return np.fft.ifft(result).real 

31 

32 @staticmethod 

33 def d2shift(img, offset, crop=False): 

34 """ 

35 Creates a FFTd version of the image and shifts it by the offset amount. 

36 

37 WARNING The shift applied is -offset, thus a positive value of offset means a shift torwards 0,0 

38 

39 ## Params 

40 - `ìmg`: An array object representing the image 

41 - `offset`: The amount the center is shifted must be float or iterable. 

42 If float the same offset will be applied to all axes, if iterable an offset for each axis must be provided 

43 """ 

44 offset = -np.array(offset) 

45 if np.all(np.array(offset) == 0): 

46 return img 

47 ft_img = np.fft.fft2(img) 

48 result = ndimage.fourier_shift(ft_img, shift=offset) 

49 result = np.fft.ifft2(result).real 

50 return result 

51 

52 

53def round_shift_for_roi(shift): 

54 """ 

55 Round shift wrapper to compute the roi of the shifted image 

56 """ 

57 return np.round(shift).astype(int)