Coverage for src/bob/bio/face/script/crop_face_106_landmarks.py: 0%

1# CLI Command that crops faces using the eyes annotations (FaceCrop) 

2########################### 

3 

4import os 

5 

6import cv2 

7import dask.bag 

8import numpy as np 

9 

10from clapper.click import ResourceOption 

11from skimage import transform as trans 

12 

13import bob.bio.face 

14import bob.io.base 

15 

16from bob.bio.face.annotator import FaceX106Landmarks 

17from bob.io.image import bob_to_opencvbgr, opencvbgr_to_bob 

18from bob.pipelines.distributed import VALID_DASK_CLIENT_STRINGS 

19 

20# Taken from here: https://github.com/JDAI-CV/FaceX-Zoo/blob/db0b087e4f4d28152e172d6c8d3767a8870733b4/face_sdk/utils/lms_trans.py 

21 

22lms5_2_lms106 = {1: 105, 2: 106, 3: 55, 4: 85, 5: 91} 

23 

24 

25def lms106_2_lms5(lms_106): 

26 lms5 = [] 

27 for cur_point_index in range(5): 

28 cur_point_id = cur_point_index + 1 

29 point_id_106 = lms5_2_lms106[cur_point_id] 

30 cur_point_index_106 = point_id_106 - 1 

31 cur_point_x = lms_106[cur_point_index_106 * 2] 

32 cur_point_y = lms_106[cur_point_index_106 * 2 + 1] 

33 lms5.append(cur_point_x) 

34 lms5.append(cur_point_y) 

35 return lms5 

36 

37 

38# Taken from here: https://github.com/JDAI-CV/FaceX-Zoo/blob/db0b087e4f4d28152e172d6c8d3767a8870733b4/face_sdk/utils/lms_trans.py 

39def estimate_norm(lmk, image_size=112): 

40 # Arcface reference points for aligment 

41 arcface_reference_lmk = np.array( 

42 [ 

43 [38.2946, 51.6963], 

44 [73.5318, 51.5014], 

45 [56.0252, 71.7366], 

46 [41.5493, 92.3655], 

47 [70.7299, 92.2041], 

48 ], 

49 dtype=np.float32, 

50 ) 

51 

52 arcface_reference_lmk = np.expand_dims(arcface_reference_lmk, axis=0) 

53 

54 ################ 

55 

56 assert lmk.shape == (5, 2) 

57 tform = trans.SimilarityTransform() 

58 lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1) 

59 min_M = [] 

60 min_index = [] 

61 min_error = float("inf") 

62 for i in np.arange(arcface_reference_lmk.shape[0]): 

63 tform.estimate(lmk, arcface_reference_lmk[i]) 

64 M = tform.params[0:2, :] 

65 results = np.dot(M, lmk_tran.T) 

66 results = results.T 

67 error = np.sum( 

68 np.sqrt(np.sum((results - arcface_reference_lmk[i]) ** 2, axis=1)) 

69 ) 

70 # print(error) 

71 if error < min_error: 

72 min_error = error 

73 min_M = M 

74 min_index = i 

75 return min_M, min_index 

76 

77 

78def faceX_cropper( 

79 files, 

80 database_path, 

81 output_path, 

82): 

83 annotator = FaceX106Landmarks() 

84 

85 image_size = 112 

86 

87 # Load 

88 for f in files: 

89 f = f.rstrip("\n") 

90 

91 output_filename = os.path.join(output_path, f) 

92 if os.path.exists(output_filename): 

93 continue 

94 

95 image = bob.io.base.load(os.path.join(database_path, f)) 

96 

97 # If it's grayscaled, expand dims 

98 if image.ndim == 2: 

99 image = np.repeat(np.expand_dims(image, 0), 3, axis=0) 

100 

101 # DEtect landmarks 

102 

103 annot = annotator.annotate(image.copy()) 

104 

105 if annot is None: 

106 print(f"Face on {f} was not detected") 

107 else: 

108 annot = annot.flatten() 

109 

110 landmarks = np.array(lms106_2_lms5(annot)) 

111 landmarks = landmarks.reshape((5, 2)) 

112 

113 M, pose_index = estimate_norm(landmarks, image_size=image_size) 

114 

115 # bob_to_opencvbgr, opencvbgr_to_bob 

116 image = bob_to_opencvbgr(image) 

117 

118 cropped_image = cv2.warpAffine( 

119 image.copy(), M, (image_size, image_size), borderValue=0.0 

120 ) 

121 

122 cropped_image = opencvbgr_to_bob(cropped_image) 

123 

124 os.makedirs(os.path.dirname(output_filename), exist_ok=True) 

125 bob.io.base.save(cropped_image, output_filename) 

126 

127 pass 

128 

129 

130import click 

131 

132 

133@click.command() 

134@click.argument("file_list") 

135@click.argument("database_path") 

136@click.argument("output_path") 

137@click.option( 

138 "--dask-client", 

139 "-l", 

140 entry_point_group="dask.client", 

141 string_exceptions=VALID_DASK_CLIENT_STRINGS, 

142 default="single-threaded", 

143 help="Dask client for the execution of the pipeline.", 

144 cls=ResourceOption, 

145) 

146def crop_faces_faceX(file_list, database_path, output_path, dask_client): 

147 files = open(file_list).readlines() 

148 

149 files = dask.bag.from_sequence(files) 

150 files.map_partitions( 

151 faceX_cropper, 

152 database_path, 

153 output_path, 

154 ).compute(scheduler=dask_client) 

155 

156 print("##############################################") 

157 print("#################### DONE ####################") 

158 print("##############################################") 

159 

160 pass 

161 

162 

163if __name__ == "__main__": 

164 crop_faces_faceX()