Coverage for src/bob/bio/face/preprocessor/HistogramEqualization.py: 91%

1#!/usr/bin/env python 

2# vim: set fileencoding=utf-8 : 

3# @author: Manuel Guenther <Manuel.Guenther@idiap.ch> 

4# @author: Tiago de Freitas Pereira <tiago.pereira@idiap.ch> 

5 

6import cv2 

7import numpy 

8import numpy as np 

9 

10from .Base import Base 

11from .utils import load_cropper 

12 

13 

14class HistogramEqualization(Base): 

15 """Crops the face (if desired) and performs histogram equalization to photometrically enhance the image. 

16 

17 Parameters 

18 ---------- 

19 

20 face_cropper : str or :py:class:`bob.bio.face.preprocessor.FaceCrop` or :py:class:`bob.bio.face.preprocessor.FaceDetect` or ``None`` 

21 The face image cropper that should be applied to the image. 

22 If ``None`` is selected, no face cropping is performed. 

23 Otherwise, the face cropper might be specified as a registered resource, a configuration file, or an instance of a preprocessor. 

24 

25 .. note:: The given class needs to contain a ``crop_face`` method. 

26 

27 kwargs 

28 Remaining keyword parameters passed to the :py:class:`Base` constructor, such as ``color_channel`` or ``dtype``. 

29 """ 

30 

31 def __init__(self, face_cropper, **kwargs): 

32 Base.__init__(self, **kwargs) 

33 

34 self.face_cropper = (face_cropper,) 

35 self.cropper = load_cropper(face_cropper) 

36 

37 def equalize_histogram(self, image): 

38 """equalize_histogram(image) -> equalized 

39 

40 Performs the histogram equalization on the given image. 

41 

42 Parameters 

43 ---------- 

44 

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

46 The image to berform histogram equalization with. 

47 The image will be transformed to type ``uint8`` before computing the histogram. 

48 

49 Returns 

50 ------- 

51 

52 equalized : 2D :py:class:`numpy.ndarray` (float) 

53 The photometrically enhanced image. 

54 """ 

55 return cv2.equalizeHist(np.round(image).astype(numpy.uint8)) 

56 

57 def transform(self, X, annotations=None): 

58 """ 

59 Aligns the given image according to the given annotations. 

60 

61 First, the desired color channel is extracted from the given image. 

62 Afterward, the face is eventually cropped using the ``face_cropper`` specified in the constructor. 

63 Then, the image is photometrically enhanced using histogram equalization. 

64 Finally, the resulting face is converted to the desired data type. 

65 

66 Parameters 

67 ---------- 

68 

69 X : 2D or 3D :py:class:`numpy.ndarray` 

70 The face image to be processed. 

71 

72 annotations : dict or ``None`` 

73 The annotations that fit to the given image. 

74 Might be ``None``, when the ``face_cropper`` is ``None`` or of type :py:class:`FaceDetect`. 

75 

76 Returns 

77 ------- 

78 

79 face : 2D :py:class:`numpy.ndarray` 

80 The cropped and photometrically enhanced face. 

81 """ 

82 

83 def _crop(image, annotations): 

84 image = self.change_color_channel(image) 

85 if self.cropper is not None: 

86 # TODO: USE THE TAG `ALLOW_ANNOTATIONS` 

87 image = ( 

88 self.cropper.transform([image]) 

89 if annotations is None 

90 else self.cropper.transform([image], [annotations]) 

91 ) 

92 image = self.equalize_histogram(image[0]) 

93 else: 

94 # Handle with the cropper is None 

95 image = self.equalize_histogram(image) 

96 

97 return self.data_type(image) 

98 

99 if annotations is None: 

100 return [_crop(data) for data in X] 

101 else: 

102 return [_crop(data, annot) for data, annot in zip(X, annotations)]