Coverage for src/susi/reduc/hot_pixels/hot_pixels_correction.py: 81%

1#!/usr/bin/env python3 

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

3""" 

4hot_pixels module provides HotPixel correction class 

5 

6@author: iglesias 

7""" 

8from typing import Iterable 

9import numpy as np 

10from scipy import ndimage 

11 

12from ...base import Logging, Config, Globals 

13from ...io import Fits 

14 

15logger = Logging.get_logger() 

16 

17 

18class HOTPXCorrector: 

19 """ 

20 Class to filter outlier pixels from the input images (they are replaced by they neithgbouring median). 

21 Outliers are those beyod np.mean(diff)+-self.std_criterion*np.std(diff), where is the difference between 

22 the input image and its smoothed version (using a median filter with window size self.median_win). 

23 The main parameters controlling the algorithm are given in _load_corr_mode. 

24 """ 

25 def __init__(self, config: Config, data: Fits): 

26 self.config = config 

27 # Expected data shape is (frames, x, y) 

28 self.data = data.data.astype(float) 

29 self.std_criterion = None 

30 self.median_win = None 

31 self.use_edges = None 

32 

33 def run(self) -> np.array: 

34 if self._load_corr_mode(): 

35 logger.debug('Running hot pixel correction') 

36 for i in range(self.data.shape[0]): 

37 self._correct_frame(i) 

38 return self.data 

39 

40 def _load_corr_mode(self) -> tuple: 

41 """ 

42 :param self.mode defines the type of image to filter. Supported modes: 

43 -"dark" 

44 -"polcal" 

45 -"slitjaw" TODO 

46 -"spectra" TODO 

47 Returns: 

48 std_criterion: Criterion to identified statistical outliers 

49 median_win: size of the median filter window [x,y] 

50 use_edges: Set to also filter the edges (border row and col). It is slower 

51 """ 

52 try: 

53 self.std_criterion = Globals.CAM_HOT_PX_MODES[self.config.cam.hot_px_mode]['std'] 

54 self.median_win = Globals.CAM_HOT_PX_MODES[self.config.cam.hot_px_mode]['win'] 

55 self.use_edges = Globals.CAM_HOT_PX_MODES[self.config.cam.hot_px_mode]['use_edges'] 

56 except KeyError: 

57 logger.warning('Value of config.cam_hot_px_mode not recognized, skipping hot px correction') 

58 return False 

59 return True 

60 

61 def _correct_frame(self, frame_no) -> np.array: 

62 """ 

63 Filters out outlier pixels in frame number frame_no of self.data 

64 

65 INPUTS 

66 :param img: numpy array with a single image [x,y] 

67 """ 

68 img = self.data[frame_no] 

69 blurred = ndimage.median_filter(img, size=self.median_win) 

70 difference = img - blurred 

71 threshold = np.mean(difference) + (self.std_criterion * np.std(difference)) 

72 # find the hot pixels, but ignore the edges 

73 hot_pixels = np.array(np.nonzero((np.abs(difference[1:-1, 1:-1]) > threshold))) + 1 

74 

75 result = np.copy(img) 

76 for y, x in zip(hot_pixels[0], hot_pixels[1]): 

77 result[y, x] = blurred[y, x] 

78 

79 if self.use_edges: 

80 height, width = np.shape(img) 

81 

82 # ## Now get the pixels on the edges (but not the corners) ## # 

83 # left and right sides 

84 for index in range(1, height - 1): 

85 # left side: 

86 med = np.median(img[index - 1:index + 2, 0:2]) 

87 diff = np.abs(img[index, 0] - med) 

88 if diff > threshold: 

89 hot_pixels = np.hstack((hot_pixels, [[index], [0]])) 

90 result[index, 0] = med 

91 

92 # right side: 

93 med = np.median(img[index - 1:index + 2, -2:]) 

94 diff = np.abs(img[index, -1] - med) 

95 if diff > threshold: 

96 hot_pixels = np.hstack((hot_pixels, [[index], [width - 1]])) 

97 result[index, -1] = med 

98 

99 # Then the top and bottom 

100 for index in range(1, width - 1): 

101 # bottom: 

102 med = np.median(img[0:2, index - 1:index + 2]) 

103 diff = np.abs(img[0, index] - med) 

104 if diff > threshold: 

105 hot_pixels = np.hstack((hot_pixels, [[0], [index]])) 

106 result[0, index] = med 

107 

108 # top: 

109 med = np.median(img[-2:, index - 1:index + 2]) 

110 diff = np.abs(img[-1, index] - med) 

111 if diff > threshold: 

112 hot_pixels = np.hstack((hot_pixels, [[height - 1], [index]])) 

113 result[-1, index] = med 

114 

115 # ## Then the corners ## # 

116 # bottom left 

117 med = np.median(img[0:2, 0:2]) 

118 diff = np.abs(img[0, 0] - med) 

119 if diff > threshold: 

120 hot_pixels = np.hstack((hot_pixels, [[0], [0]])) 

121 result[0, 0] = med 

122 

123 # bottom right 

124 med = np.median(img[0:2, -2:]) 

125 diff = np.abs(img[0, -1] - med) 

126 if diff > threshold: 

127 hot_pixels = np.hstack((hot_pixels, [[0], [width - 1]])) 

128 result[0, -1] = med 

129 

130 # top left 

131 med = np.median(img[-2:, 0:2]) 

132 diff = np.abs(img[-1, 0] - med) 

133 if diff > threshold: 

134 hot_pixels = np.hstack((hot_pixels, [[height - 1], [0]])) 

135 result[-1, 0] = med 

136 

137 # top right 

138 med = np.median(img[-2:, -2:]) 

139 diff = np.abs(img[-1, -1] - med) 

140 if diff > threshold: 

141 hot_pixels = np.hstack((hot_pixels, [[height - 1], [width - 1]])) 

142 result[-1, -1] = med 

143 

144 self.data[frame_no] = result