-
Notifications
You must be signed in to change notification settings - Fork 7
/
downsample_kittisf.py
66 lines (53 loc) · 2.22 KB
/
downsample_kittisf.py
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
66
import sys
import pathlib
root_dir = pathlib.Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(root_dir))
import os
import os.path as osp
import tqdm
import argparse
import numpy as np
from utils.data_util import fps_downsample
parser = argparse.ArgumentParser()
parser.add_argument('data_root', type=str, help='Root path for the dataset')
parser.add_argument('--save_root', type=str, help='Path to save the downsampled dataset')
parser.add_argument('--predflow_path', type=str, default=None, help='Path to load pre-saved flow predictions')
args = parser.parse_args()
data_root = args.data_root
save_root = args.save_root
predflow_path = args.predflow_path
n_sample_point = 8192
SAVE_DIR = osp.join(save_root, 'data')
os.makedirs(SAVE_DIR, exist_ok=True)
if predflow_path is not None:
SAVE_PREDFLOW_DIR = osp.join(save_root, 'flow_preds', predflow_path)
os.makedirs(SAVE_PREDFLOW_DIR, exist_ok=True)
# Setup the dataset
from datasets.dataset_kittisf import KITTISceneFlowDataset
mapping_path = 'data_prepare/kittisf/splits/all.txt'
dataset = KITTISceneFlowDataset(data_root=data_root,
mapping_path=mapping_path,
downsampled=False,
predflow_path=predflow_path)
n_scenes = len(dataset)
pbar = tqdm.tqdm(total=n_scenes)
for sid in range(n_scenes):
pcs, segms, flows = dataset[sid]
pc_org, segm_org, flow_org = pcs[0], segms[0], flows[0]
fps_idx = fps_downsample(pc_org, n_sample_point=n_sample_point)
pc = pc_org[fps_idx]
segm = segm_org[fps_idx]
flow = flow_org[fps_idx]
idx, view_sel_idx = sid // 2, sid % 2
data_id = dataset.data_ids[idx]
save_path = osp.join(SAVE_DIR, data_id)
os.makedirs(save_path, exist_ok=True)
np.save(osp.join(save_path, 'pc%d.npy'%(view_sel_idx + 1)), pc)
np.save(osp.join(save_path, 'segm%d.npy'%(view_sel_idx + 1)), segm)
if predflow_path is not None:
save_predflow_path = osp.join(SAVE_PREDFLOW_DIR, data_id)
os.makedirs(save_predflow_path, exist_ok=True)
np.save(osp.join(save_predflow_path, 'flow%d.npy'%(view_sel_idx + 1)), flow)
else:
np.save(osp.join(save_path, 'flow%d.npy'%(view_sel_idx + 1)), flow)
pbar.update()