Spaces:
Running
Running
File size: 2,435 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
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 |
import os
import glob
from re import split
from tqdm import tqdm
from multiprocessing import Pool
from functools import partial
scannet_dir = "/root/data/ScanNet-v2-1.0.0/data/raw"
dump_dir = "/root/data/scannet_dump"
num_process = 32
def extract(seq, scannet_dir, split, dump_dir):
assert split == "train" or split == "test"
if not os.path.exists(os.path.join(dump_dir, split, seq)):
os.mkdir(os.path.join(dump_dir, split, seq))
cmd = (
"python reader.py --filename "
+ os.path.join(
scannet_dir,
"scans" if split == "train" else "scans_test",
seq,
seq + ".sens",
)
+ " --output_path "
+ os.path.join(dump_dir, split, seq)
+ " --export_depth_images --export_color_images --export_poses --export_intrinsics"
)
os.system(cmd)
if __name__ == "__main__":
if not os.path.exists(dump_dir):
os.mkdir(dump_dir)
os.mkdir(os.path.join(dump_dir, "train"))
os.mkdir(os.path.join(dump_dir, "test"))
train_seq_list = [
seq.split("/")[-1]
for seq in glob.glob(os.path.join(scannet_dir, "scans", "scene*"))
]
test_seq_list = [
seq.split("/")[-1]
for seq in glob.glob(os.path.join(scannet_dir, "scans_test", "scene*"))
]
extract_train = partial(
extract, scannet_dir=scannet_dir, split="train", dump_dir=dump_dir
)
extract_test = partial(
extract, scannet_dir=scannet_dir, split="test", dump_dir=dump_dir
)
num_train_iter = (
len(train_seq_list) // num_process
if len(train_seq_list) % num_process == 0
else len(train_seq_list) // num_process + 1
)
num_test_iter = (
len(test_seq_list) // num_process
if len(test_seq_list) % num_process == 0
else len(test_seq_list) // num_process + 1
)
pool = Pool(num_process)
for index in tqdm(range(num_train_iter)):
seq_list = train_seq_list[
index * num_process : min((index + 1) * num_process, len(train_seq_list))
]
pool.map(extract_train, seq_list)
pool.close()
pool.join()
pool = Pool(num_process)
for index in tqdm(range(num_test_iter)):
seq_list = test_seq_list[
index * num_process : min((index + 1) * num_process, len(test_seq_list))
]
pool.map(extract_test, seq_list)
pool.close()
pool.join()
|