File size: 2,324 Bytes
4f6b78d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import sys
import torch
sys.path.append('.')
import os
import numpy as np
import glob
from tqdm import tqdm
dataset_location = '../data/point_odyssey'
# print(dataset_location)
for dset in ["train", "test", "sample"]:
sequences = []
subdir = os.path.join(dataset_location, dset)
for seq in glob.glob(os.path.join(subdir, "*/")):
sequences.append(seq)
# sequences = sorted(sequences)
squences = sorted(sequences)
print('found %d unique videos in %s (dset=%s)' % (len(sequences), dataset_location, dset))
## load trajectories
print('loading trajectories...')
for seq in sequences:
# print('seq', seq)
# if os.path.exists(os.path.join(seq, 'trajs_2d')):
# print('skipping', seq)
# continue
info_path = os.path.join(seq, 'info.npz')
info = np.load(info_path, allow_pickle=True)
trajs_3d_shape = info['trajs_3d'].astype(np.float32)
if len(trajs_3d_shape):
print('processing', seq)
rgb_path = os.path.join(seq, 'rgbs')
info_path = os.path.join(seq, 'info.npz')
annotations_path = os.path.join(seq, 'anno.npz')
trajs_3d_path = os.path.join(seq, 'trajs_3d')
trajs_2d_path = os.path.join(seq, 'trajs_2d')
os.makedirs(trajs_3d_path, exist_ok=True)
os.makedirs(trajs_2d_path, exist_ok=True)
info = np.load(info_path, allow_pickle=True)
trajs_3d_shape = info['trajs_3d']
anno = np.load(annotations_path, allow_pickle=True)
keys = {'trajs_2d': 'traj_2d', 'trajs_3d': 'traj_3d', 'valids': 'valid', 'visibs': 'visib', 'intrinsics': 'intrinsic', 'extrinsics': 'extrinsic'}
if len(trajs_3d_shape) == 0:
print(anno['trajs_3d'])
print('skipping', seq)
continue
tensors = {key: torch.tensor(anno[key]).cuda() for key in keys}
for t in tqdm(range(trajs_3d_shape[0])):
for key, item_key in keys.items():
path = os.path.join(seq, key)
os.makedirs(path, exist_ok=True)
filename = os.path.join(path, f'{item_key}_{t:05d}.npy')
np.save(filename, tensors[key][t].cpu().numpy())
|