Spaces:
Runtime error
Runtime error
File size: 2,904 Bytes
1bb1365 |
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 |
#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Script to export the list of scenes for habitat (after having rendered them).
# Usage:
# python3 datasets_preprocess/preprocess_co3d.py --root data/habitat_processed
# --------------------------------------------------------
import os
from collections import defaultdict
import numpy as np
from tqdm import tqdm
def find_all_scenes(habitat_root, n_scenes=[100000]):
np.random.seed(777)
try:
fpath = os.path.join(habitat_root, f"Habitat_all_scenes.txt")
list_subscenes = open(fpath).read().splitlines()
except IOError:
if input("parsing sub-folders to find scenes? (y/n) ") != "y":
return
list_subscenes = []
for root, dirs, files in tqdm(os.walk(habitat_root)):
for f in files:
if not f.endswith("_1_depth.exr"):
continue
scene = os.path.join(
os.path.relpath(root, habitat_root), f.replace("_1_depth.exr", "")
)
if hash(scene) % 1000 == 0:
print("... adding", scene)
list_subscenes.append(scene)
with open(fpath, "w") as f:
f.write("\n".join(list_subscenes))
print(f">> wrote {fpath}")
print(f"Loaded {len(list_subscenes)} sub-scenes")
# separate scenes
list_scenes = defaultdict(list)
for scene in list_subscenes:
scene, id = os.path.split(scene)
list_scenes[scene].append(id)
list_scenes = list(list_scenes.items())
print(f"from {len(list_scenes)} scenes in total")
np.random.shuffle(list_scenes)
train_scenes = list_scenes[len(list_scenes) // 10 :]
val_scenes = list_scenes[: len(list_scenes) // 10]
def write_scene_list(scenes, n, fpath):
sub_scenes = [os.path.join(scene, id) for scene, ids in scenes for id in ids]
np.random.shuffle(sub_scenes)
if len(sub_scenes) < n:
return
with open(fpath, "w") as f:
f.write("\n".join(sub_scenes[:n]))
print(f">> wrote {fpath}")
for n in n_scenes:
write_scene_list(
train_scenes, n, os.path.join(habitat_root, f"Habitat_{n}_scenes_train.txt")
)
write_scene_list(
val_scenes,
n // 10,
os.path.join(habitat_root, f"Habitat_{n//10}_scenes_val.txt"),
)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--root", required=True)
parser.add_argument(
"--n_scenes", nargs="+", default=[1_000, 10_000, 100_000, 1_000_000], type=int
)
args = parser.parse_args()
find_all_scenes(args.root, args.n_scenes)
|