Spaces:
Running
Running
import os | |
import json | |
import tables | |
from tqdm import tqdm | |
import numpy as np | |
def read_all_imgs(base_dir): | |
all_imgs = [] | |
for cur, dirs, files in os.walk(base_dir): | |
if 'imgs' in cur: | |
all_imgs += [os.path.join(cur, f) for f in files] | |
all_imgs.sort() | |
return all_imgs | |
def filter_semantic_depth(imgs): | |
valid_imgs = [] | |
for item in tqdm(imgs): | |
f_name = os.path.splitext(os.path.basename(item))[0] + '.h5' | |
depth_dir = os.path.abspath(os.path.join(os.path.dirname(item), '../depths')) | |
depth_path = os.path.join(depth_dir, f_name) | |
depth_h5 = tables.open_file(depth_path, mode='r') | |
_depth = np.array(depth_h5.root.depth) | |
if _depth.min() >= 0: | |
prefix = os.path.abspath(os.path.join(item, '../../../../')) + '/' | |
rel_image_path = item.replace(prefix, '') | |
valid_imgs.append(rel_image_path) | |
depth_h5.close() | |
valid_imgs.sort() | |
return valid_imgs | |
if __name__ == "__main__": | |
MegaDepth_v1 = '/media/jiangwei/data_ssd/MegaDepth_v1/' | |
assert os.path.isdir(MegaDepth_v1), 'Change to your local path' | |
all_imgs = read_all_imgs(MegaDepth_v1) | |
valid_imgs = filter_semantic_depth(all_imgs) | |
with open('megadepth_valid_list.json', 'w') as outfile: | |
json.dump(valid_imgs, outfile, indent=4) | |