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