Coverage for src/spectroflat/fitting/line_locator.py: 100%

1import numpy as np 

2from scipy import signal as sig 

3 

4from ..base import Logging 

5 

6log = Logging.get_logger() 

7 

8 

9class AutoLineLocator: 

10 

11 def __init__(self, image: np.array, line_distance: int = 100, sigma: float = 0.6, line_prominence=0): 

12 self.img = image 

13 self.sigma = sigma 

14 self.line_distance = line_distance 

15 self.line_prominence = line_prominence 

16 

17 def detect_centers(self) -> np.array: 

18 self._normalize() 

19 return self._detect_peaks() 

20 

21 def _normalize(self): 

22 self.img = self.img / self.img.std() 

23 self.img = self.img - self.img.min() 

24 

25 def _detect_peaks(self) -> np.array: 

26 row_means = self.img.mean(axis=0) 

27 h = np.mean(row_means) + self.sigma * np.std(row_means) 

28 peaks, _ = sig.find_peaks(row_means, height=h, distance=self.line_distance, prominence=self.line_prominence) 

29 log.info('Detected %s lines at:\n%s', len(peaks), peaks) 

30 return peaks