File size: 480 Bytes
bc3753a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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