Coverage for src/spectroflat/shift/img_rotation.py: 84%

1import numpy as np 

2 

3from skimage.feature import canny 

4from skimage.transform import probabilistic_hough_line, rotate 

5 

6from ..base.exceptions import InsufficientDataException 

7 

8DEG2RAD = np.pi / 180 

9 

10 

11class RotationAnalysis: 

12 """ 

13 This class analyzes the global rotation of a given image. 

14 

15 First it detects edges via the canny algorithm, and then it uses scikit 

16 implementation of [Hough Transform](https://en.wikipedia.org/wiki/Hough_transform) to detect 

17 straight lines. Finally, the angle of all lines is computed and the median of all angles is returned. 

18 """ 

19 

20 @staticmethod 

21 def detect_horizontal_rotation(img: np.array) -> float: 

22 return RotationAnalysis(img.T).run().angle * -1 

23 

24 @staticmethod 

25 def detect_vertical_rotation(img: np.array) -> float: 

26 return RotationAnalysis(img).run().angle 

27 

28 def __init__(self, img: np.array, deg_range: float = 4): 

29 #: The image to analyse 

30 self._orig = img 

31 #: Max rotation in deg 

32 self._deg_range = deg_range 

33 #: The detected rotation angle in degree 

34 self.angle: float = 0.0 

35 

36 def run(self): 

37 """Start the detection algorithm""" 

38 lines = self._detect_lines() 

39 if len(lines) == 0: 

40 raise InsufficientDataException('No lines detected. Cannot determine image rotation.') 

41 

42 self._detect_rotation(lines) 

43 return self 

44 

45 def _detect_lines(self) -> list: 

46 theta = np.linspace(-self._deg_range * DEG2RAD, self._deg_range * DEG2RAD) 

47 edges = canny(self._orig, sigma=1.5) 

48 return probabilistic_hough_line(edges, threshold=50, line_length=150, line_gap=10, theta=theta) 

49 

50 def _detect_rotation(self, lines: list): 

51 angles = [np.degrees(np.arctan2(y2 - y1, x2 - x1)) for (x1, y1), (x2, y2) in lines] 

52 # noinspection PyTypeChecker 

53 self.angle = RotationAnalysis.minimize(np.median(angles)) 

54 

55 @staticmethod 

56 def minimize(angle: float) -> float: 

57 angle = angle % 90 

58 if angle > 45: 

59 angle -= 90 

60 if angle < -45: 

61 angle += 90 

62 return angle 

63 

64 

65class RotationCorrection: 

66 

67 def __init__(self, img: np.array, angle: float): 

68 self._img = img 

69 self._angle = angle 

70 

71 def _cut_shape(self) -> tuple: 

72 rad_angle = np.abs(self._angle) * DEG2RAD 

73 cut0 = int(np.ceil(np.tan(rad_angle) * self._img.shape[1] / 2)) 

74 cut1 = int(np.ceil(np.tan(rad_angle) * self._img.shape[0] / 2)) 

75 return slice(cut0, self._img.shape[0] - cut0), slice(cut1, self._img.shape[1] - cut1) 

76 

77 def bicubic(self, cut_shape: bool = False, pad_mean: bool = True) -> np.array: 

78 """ 

79 Rotate the given image by the given angle with a bicubic algorithm 

80 ### Params 

81 - cut_shape: If True image will be cropped to square 

82 - pad_mean: If True the black borders will be filled with the images mean value instead. 

83 

84 ### Returns 

85 The rotated image 

86 """ 

87 if self._angle == 0: 

88 return self._img 

89 img = rotate(self._img, self._angle, order=3) 

90 if cut_shape: 

91 return img[self._cut_shape()] 

92 if pad_mean: 

93 inner = self._cut_shape() 

94 temp = np.ones(self._img.shape) * img[inner].mean() 

95 temp[inner] = img[inner] 

96 img = temp 

97 return img