Coverage for src/bob/bio/face/preprocessor/Base.py: 74%

1import numpy 

2 

3from sklearn.base import BaseEstimator, TransformerMixin 

4 

5from bob.bio.face.color import gray_to_rgb, rgb_to_gray 

6 

7 

8def change_color_channel(image, color_channel): 

9 if image.ndim == 2: 

10 if color_channel == "rgb": 

11 return gray_to_rgb(image) 

12 if color_channel != "gray": 

13 raise ValueError( 

14 "There is no rule to extract a " 

15 + color_channel 

16 + " image from a gray level image!" 

17 ) 

18 return image 

19 

20 if color_channel == "rgb": 

21 return image 

22 if color_channel == "bgr": 

23 return image[[2, 1, 0], ...] 

24 if color_channel == "gray": 

25 return rgb_to_gray(image) 

26 if color_channel == "red": 

27 return image[0, :, :] 

28 if color_channel == "green": 

29 return image[1, :, :] 

30 if color_channel == "blue": 

31 return image[2, :, :] 

32 

33 raise ValueError( 

34 "The image channel '%s' is not known or not yet implemented", 

35 color_channel, 

36 ) 

37 

38 

39class Base(TransformerMixin, BaseEstimator): 

40 """Performs color space adaptations and data type corrections for the given 

41 image. 

42 

43 **Parameters:** 

44 

45 dtype : :py:class:`numpy.dtype` or convertible or ``None`` 

46 The data type that the resulting image will have. 

47 

48 color_channel : one of ``('gray', 'red', 'gren', 'blue', 'rgb')`` 

49 The specific color channel, which should be extracted from the image. 

50 """ 

51 

52 def __init__(self, dtype=None, color_channel="gray", **kwargs): 

53 self.color_channel = color_channel 

54 self.dtype = dtype 

55 

56 @property 

57 def channel(self): 

58 return self.color_channel 

59 

60 def _more_tags(self): 

61 return {"requires_fit": False} 

62 

63 def fit(self, X, y=None): 

64 return self 

65 

66 def change_color_channel(self, image): 

67 """color_channel(image) -> channel 

68 

69 Returns the channel of the given image, which was selected in the 

70 constructor. Currently, gray, red, green and blue channels are supported. 

71 

72 **Parameters:** 

73 

74 image : 2D or 3D :py:class:`numpy.ndarray` 

75 The image to get the specified channel from. 

76 

77 **Returns:** 

78 

79 channel : 2D or 3D :py:class:`numpy.ndarray` 

80 The extracted color channel. 

81 """ 

82 

83 return change_color_channel(image, self.color_channel) 

84 

85 def data_type(self, image): 

86 """ 

87 Converts the given image into the data type specified in the constructor of 

88 this class. If no data type was specified, or the ``image`` is ``None``, no 

89 conversion is performed. 

90 

91 Parameters 

92 ---------- 

93 

94 image : 2D or 3D :py:class:`numpy.ndarray` 

95 The image to convert. 

96 

97 Returns 

98 ------- 

99 

100 image : 2D or 3D :py:class:`numpy.ndarray` 

101 The image converted to the desired data type, if any. 

102 """ 

103 if self.dtype is not None and image is not None: 

104 image = image.astype(self.dtype) 

105 return image 

106 

107 def transform(self, images, annotations=None): 

108 """ 

109 Extracts the desired color channel and converts to the desired data type. 

110 

111 Parameters 

112 ---------- 

113 

114 image : 2D or 3D :py:class:`numpy.ndarray` 

115 The image to preprocess. 

116 

117 annotations : any 

118 Ignored. 

119 

120 Returns 

121 ------- 

122 

123 image : 2D :py:class:`numpy.ndarray` 

124 The image converted converted to the desired color channel and type. 

125 """ 

126 return [self._transform_one_image(img) for img in images] 

127 

128 def _transform_one_image(self, image): 

129 assert isinstance(image, numpy.ndarray) and image.ndim in (2, 3) 

130 # convert to grayscale 

131 image = self.change_color_channel(image) 

132 return self.data_type(image)