Coverage for src/susi/analyse/contrast.py: 36%

1""" 

2Compute contrast of a series of frames. 

3""" 

4from __future__ import annotations 

5 

6import glob 

7import logging 

8import numpy as np 

9 

10from ..io import FitsBatch 

11from ..base import Config, InsufficientDataException 

12from ..utils import MP 

13 

14log = logging.getLogger('SUSI') 

15 

16 

17class ContrastAnalyser: 

18 """ 

19 Compute contrast of a series of frames. 

20 """ 

21 def __init__(self, config: Config, workers: int = None, upscale: int = 1): 

22 self.batch = FitsBatch(config=config, slices=config.cam.data_shape) 

23 #: The config to apply (i.e. defines the frame set to analyse) 

24 self.config = config 

25 #: The number of parallel processes to use 

26 self.workers = config.base.workers 

27 #: The upscale factor to analyse sub-pixel shifts(default 1) 

28 self.upscale = upscale 

29 #: Dark image 

30 self.dark_file = None 

31 #: The result 

32 self.result = [] 

33 

34 def with_dark_correction(self, dark_file): 

35 """ 

36 Chainable configuration method to activate dark img correction 

37 

38 :param file: [String] the path of the darks file to use 

39 :return: self 

40 """ 

41 self.dark_file = dark_file 

42 return self 

43 

44 def run(self) -> ContrastAnalyser: 

45 self.__load_data() 

46 self.__apply_dark() 

47 self.__analyse() 

48 return self 

49 

50 def __load_data(self): 

51 file_list = self.config.find_files() 

52 log.debug('Loading area %s of all files ...', self.config.cam.data_shape) 

53 self.batch.load(file_list, workers=min(len(file_list), 100), sort_by=self.config.base.timestamp_field) 

54 

55 def __apply_dark(self): 

56 if self.dark_file is not None: 

57 self.batch.apply_dark(self.dark_file) 

58 

59 def __analyse(self): 

60 frames = self.batch.data_array() 

61 args = [{'idx': i, 'frame': frames[i]} for i in range(0, len(frames))] 

62 log.debug('Compute contrast') 

63 result = dict(MP.simultaneous(_compute_contrast, args, workers=self.workers)) 

64 log.debug('Analysis done, collecting results...') 

65 for i in range(len(result.keys())): 

66 self.result.append(result[i]) 

67 

68 

69def _compute_contrast(data: dict) -> tuple: 

70 """ 

71 :param data: a dictionary with 

72 - 'idx' the index of the frame 

73 - 'frame' the frame to correlate 

74 

75 :return: A tuple where element 0 is the idx and element 1 is the contrast value 

76 """ 

77 if data['idx'] % 200 == 0: 

78 log.debug('\tProcessing Frame %s', data['idx']) 

79 contrast = np.std(data['frame']) / np.mean(data['frame']) 

80 return data['idx'], contrast