NeRF/data_utils/face_tracking/data_loader.py

import os
import torch
import numpy as np


def load_dir(path, start, end):
    lmss = []
    imgs_paths = []
    for i in range(start, end):
        if os.path.isfile(os.path.join(path, str(i) + ".lms")):
            lms = np.loadtxt(os.path.join(path, str(i) + ".lms"), dtype=np.float32)
            lmss.append(lms)
            imgs_paths.append(os.path.join(path, str(i) + ".jpg"))
    lmss = np.stack(lmss)
    lmss = torch.as_tensor(lmss).cuda()
    return lmss, imgs_paths