import os
import sys

import numpy as np
from tqdm import tqdm

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from utils.utils import mask_faces

def generate_faces(folder):
    frames = os.listdir(os.path.join(folder, 'mesh_mask'))
    F = np.load('custom_data/smpl_faces.npy')
    os.makedirs(os.path.join(folder, 'faces'), exist_ok=True)

    for frame in tqdm(frames):
        mask = np.load(os.path.join(folder, 'mesh_mask', frame))
        faces = mask_faces(F, mask)
        np.save(os.path.join(folder, 'faces', frame), faces)


if __name__ == '__main__':
    # folder = '/home/workspace/local-data/train'
    folder = '/home/workspace/local-data/valid'

    generate_faces(folder)