from face3d.morphable_model import MorphabelModel
from face_alignment

self.bfm = MorphabelModel('/home/alex/Documents/stereo/BFM.mat')
self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2halfD, flip_input=False)


def compute_face_description(self, images, lite_stereo=False):
        def fit(landmarks):
            h, w = images[0].shape[0], images[0].shape[1]
            if landmarks and len(landmarks)==1 :
                x=np.zeros((landmarks[0].shape[0],2))
                x[:,:] = landmarks[0][:, 0:2]  # 2d keypoint, which can be detected from image
                x[:, 0] -= w / 2
                x[:, 1] -= h / 2
                x[:, 1] = -x[:, 1]
                X_ind = self.bfm.kpt_ind
                fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t = self.bfm.fit(x, X_ind, max_iter=3, isShow=False)
                return fitted_angles
            else:
                if landmarks:
                    print("Detected landmarks for %d faces, on this stage only 1 should be selected" % len(landmarks))
                return None
        #print(self.detected_faces)
        start_time = time.time()
        self.landmarks = [self.fa.get_landmarks_from_image(img[...,::-1])
                          for img in images]
        print('       descriptors found', time.time() - start_time)

        if not lite_stereo:
            angles = [fit(land) for land in self.landmarks]
        else:
            angles = [fit(land) if i < len(images)//2 else [0,0,0] for i, land in enumerate(self.landmarks)]

        print('       angles calculated', time.time() - start_time)

        return angles