17 lines
480 B
Python
17 lines
480 B
Python
|
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
|