File size: 5,215 Bytes
f53b39e |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import os
import glob
from tqdm import tqdm
# Define the merged dataset metadata dictionary
dataset_metadata = {
'davis': {
'img_path': "./data/davis/DAVIS/JPEGImages/480p",
'mask_path': "./data/davis/DAVIS/Annotations/480p",
'dir_path_func': lambda img_path, seq: os.path.join(img_path, seq),
'gt_traj_func': lambda img_path, anno_path, seq: None,
'traj_format': None,
'seq_list': None,
'full_seq': True,
'mask_path_seq_func': lambda mask_path, seq: os.path.join(mask_path, seq),
'skip_condition': None,
'process_func': None, # Not used in mono depth estimation
},
'test': {
'img_path': "./data/test_set",
'mask_path': None,
'dir_path_func': lambda img_path, seq: os.path.join(img_path, seq),
'gt_traj_func': lambda img_path, anno_path, seq: None,
'traj_format': None,
'seq_list': None,
'full_seq': True,
'mask_path_seq_func': lambda mask_path, seq: None,
'skip_condition': None,
'process_func': None, # Not used in mono depth estimation
},
'bonn': {
'img_path': "./data/bonn/rgbd_bonn_dataset",
'mask_path': None,
'dir_path_func': lambda img_path, seq: os.path.join(img_path, f'rgbd_bonn_{seq}', 'rgb_110'),
'gt_traj_func': lambda img_path, anno_path, seq: os.path.join(img_path, f'rgbd_bonn_{seq}', 'groundtruth_110.txt'),
'traj_format': 'tum',
'seq_list': ["balloon2", "crowd2", "crowd3", "person_tracking2", "synchronous"],
'full_seq': False,
'mask_path_seq_func': lambda mask_path, seq: None,
'skip_condition': None,
'process_func': lambda args, img_path: process_bonn(args, img_path),
},
'tum': {
'img_path': "./data/tum",
'mask_path': None,
'dir_path_func': lambda img_path, seq: os.path.join(img_path, seq, 'rgb_50'),
'gt_traj_func': lambda img_path, anno_path, seq: os.path.join(img_path, seq, 'groundtruth_50.txt'),
'traj_format': 'tum',
'seq_list': None,
'full_seq': True,
'mask_path_seq_func': lambda mask_path, seq: None,
'skip_condition': None,
'process_func': None,
},
'sintel': {
'img_path': "./data/MPI-Sintel/MPI-Sintel-training_images/training/final",
'anno_path': "./data/MPI-Sintel/MPI-Sintel-depth-training/training/camdata_left",
'mask_path': "./data/MPI-Sintel/MPI-Sintel-depth-training/training/dynamic_label_perfect/",
'dir_path_func': lambda img_path, seq: os.path.join(img_path, seq),
'gt_traj_func': lambda img_path, anno_path, seq: os.path.join(anno_path, seq),
'traj_format': None,
'seq_list': ["alley_2", "ambush_4", "ambush_5", "ambush_6", "cave_2", "cave_4", "market_2",
"market_5", "market_6", "shaman_3", "sleeping_1", "sleeping_2", "temple_2", "temple_3"],
'full_seq': False,
'mask_path_seq_func': lambda mask_path, seq: None,
'skip_condition': None,
'process_func': lambda args, img_path: process_sintel(args, img_path),
},
}
# Define processing functions for each dataset
def process_kitti(args, img_path):
for dir in tqdm(sorted(glob.glob(f'{img_path}/*'))):
filelist = sorted(glob.glob(f'{dir}/*.png'))
save_dir = f'{args.output_dir}/{os.path.basename(dir)}'
yield filelist, save_dir
def process_bonn(args, img_path):
if args.full_seq:
for dir in tqdm(sorted(glob.glob(f'{img_path}/*/'))):
filelist = sorted(glob.glob(f'{dir}/rgb/*.png'))
save_dir = f'{args.output_dir}/{os.path.basename(os.path.dirname(dir))}'
yield filelist, save_dir
else:
seq_list = ["balloon2", "crowd2", "crowd3", "person_tracking2", "synchronous"] if args.seq_list is None else args.seq_list
for seq in tqdm(seq_list):
filelist = sorted(glob.glob(f'{img_path}/rgbd_bonn_{seq}/rgb_110/*.png'))
save_dir = f'{args.output_dir}/{seq}'
yield filelist, save_dir
def process_nyu(args, img_path):
filelist = sorted(glob.glob(f'{img_path}/*.png'))
save_dir = f'{args.output_dir}'
yield filelist, save_dir
def process_scannet(args, img_path):
seq_list = sorted(glob.glob(f'{img_path}/*'))
for seq in tqdm(seq_list):
filelist = sorted(glob.glob(f'{seq}/color_90/*.jpg'))
save_dir = f'{args.output_dir}/{os.path.basename(seq)}'
yield filelist, save_dir
def process_sintel(args, img_path):
if args.full_seq:
for dir in tqdm(sorted(glob.glob(f'{img_path}/*/'))):
filelist = sorted(glob.glob(f'{dir}/*.png'))
save_dir = f'{args.output_dir}/{os.path.basename(os.path.dirname(dir))}'
yield filelist, save_dir
else:
seq_list = ["alley_2", "ambush_4", "ambush_5", "ambush_6", "cave_2", "cave_4", "market_2",
"market_5", "market_6", "shaman_3", "sleeping_1", "sleeping_2", "temple_2", "temple_3"]
for seq in tqdm(seq_list):
filelist = sorted(glob.glob(f'{img_path}/{seq}/*.png'))
save_dir = f'{args.output_dir}/{seq}'
yield filelist, save_dir |