Coverage for src/susi/utils/rebin.py: 93%

1import numpy as np 

2 

3 

4def upsample(img: np.array, factor: float = None, target_shape: tuple = None) -> np.array: 

5 """ 

6 Upsamples a 2D image using FFT 

7 

8 ## Params 

9 - img: The 2D image to rescale 

10 - factor: factor to enlarge the image to (Only needed if no target_shape is given). 

11 - target_shape: The targeted shape of the new image as (y-dim, x-dim) tuple. 

12 

13 ## Returns 

14 the up-scaled image. 

15 """ 

16 if factor is None and target_shape is None: 

17 raise AttributeError('factor or target_shape argument must be given') 

18 

19 fft = np.fft.fftshift(np.fft.fft2(img)) 

20 if factor is not None: 

21 x = int(np.round(fft.shape[-1] * factor / 2)) 

22 y = int(np.round(fft.shape[-2] * factor / 2)) 

23 if target_shape is not None: 

24 x = (target_shape[-1] - fft.shape[-1]) // 2 

25 y = (target_shape[-2] - fft.shape[-2]) // 2 

26 fr = np.pad(fft.real, [[x, x], [y, y]], mode='constant') 

27 fi = np.pad(fft.imag, [[x, x], [y, y]], mode='constant') 

28 

29 fft_hires = np.fft.ifftshift(fr + 1j * fi) 

30 return np.array(np.fft.ifft2(fft_hires).real)