import torch
import os
import sys
import tqdm
import glob
import numpy as np
import cv2
import face_alignment
from skimage import io
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--data_source', type=str, default='./data/input')
args = parser.parse_args()
DATA_SOURCE = args.data_source

device = torch.device('cuda:0')
fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.THREE_D, flip_input=False, face_detector='blazeface')

# DATA_SOURCE = '../../data/face_data'
# DATA_SOURCE = '../../data/face_data'
data_folder = os.path.join(DATA_SOURCE, 'images')
frame_folders = sorted(glob.glob(data_folder + '/*'))
for frame_folder in tqdm.tqdm(frame_folders):
    if 'background' in frame_folder:
        continue
    image_paths = glob.glob(frame_folder + '/image_*')
    images = np.stack([io.imread(image_path) for image_path in image_paths])
    images = torch.from_numpy(images).float().permute(0, 3, 1, 2).to(device)
    results = fa.get_landmarks_from_batch(images, return_landmark_score=True)
    for i in range(len(results[0])):
        if results[1][i] is None:
            results[0][i] = np.zeros([68, 3], dtype=np.float32)
            results[1][i] = [np.zeros([68], dtype=np.float32)]
        if len(results[1][i]) > 1:
            total_score = 0.0
            for j in range(len(results[1][i])):
                if np.sum(results[1][i][j]) > total_score:
                    total_score = np.sum(results[1][i][j])
                    landmarks_i = results[0][i][j*68:(j+1)*68]
                    scores_i = results[1][i][j:j+1]
            results[0][i] = landmarks_i
            results[1][i] = scores_i
        
    landmarks = np.concatenate([np.stack(results[0])[:, :, :2], np.stack(results[1]).transpose(0, 2, 1)], -1)
    i = 0
    for image_path in image_paths:
        landmarks_path = image_path.replace('image_', 'landmarks_').replace('.jpg', '.npy')
        np.save(landmarks_path, landmarks[i])
        i += 1