File size: 10,650 Bytes
55e58d1 |
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 116 117 118 119 120 121 122 123 124 125 126 |
import os
from tqdm import tqdm
import time
from utils.config import get_args
from pathlib import Path
import pickle
def execute_commands(commands_list, command_type, process_num):
print('====> Start', command_type)
from multiprocessing import Pool
pool = Pool(process_num)
for _ in tqdm(pool.imap_unordered(os.system, commands_list), total=len(commands_list)):
pass
pool.close()
pool.join()
pool.terminate()
print('====> Finish', command_type)
def get_seq_name_list(dataset):
if dataset == 'scannet':
file_path = 'splits/scannet.txt'
elif dataset == 'scannetpp':
file_path = 'splits/scannetpp.txt'
elif dataset == 'scannetpp_dust3r_posed':
file_path = 'splits/scannetpp.txt'
elif dataset == 'scannetpp_dust3r_filtered_depth':
file_path = 'splits/scannetpp.txt'
elif dataset == 'scannetpp_mapanything_posed':
file_path = 'splits/scannetpp.txt'
elif dataset == 'scannetpp_dust3r_unposed':
file_path = 'splits/scannetpp.txt'
elif dataset == 'matterport3d':
file_path = 'splits/matterport3d.txt'
with open(file_path, 'r') as f:
seq_name_list = f.readlines()
seq_name_list = [seq_name.strip() for seq_name in seq_name_list]
return seq_name_list
def parallel_compute(general_command, command_name, resource_type, cuda_list, seq_name_list):
cuda_num = len(cuda_list)
if resource_type == 'cuda':
commands = []
for i, cuda_id in enumerate(cuda_list):
process_seq_name = seq_name_list[i::cuda_num]
if len(process_seq_name) == 0:
continue
process_seq_name = '+'.join(process_seq_name)
command = f'CUDA_VISIBLE_DEVICES={cuda_id} {general_command % process_seq_name}'
commands.append(command)
execute_commands(commands, command_name, cuda_num)
elif resource_type == 'cpu':
commands = []
for seq_name in seq_name_list:
commands.append(f'{general_command} --seq_name {seq_name}')
execute_commands(commands, command_name, cuda_num)
def get_label_text_feature(cuda_id):
label_text_feature_path = 'data/text_features/matterport3d.npy'
if os.path.exists(label_text_feature_path):
return
command = f'CUDA_VISIBLE_DEVICES={cuda_id} python -m semantics.extract_label_featrues'
os.system(command)
def main(args):
CUDA_LIST = args.devices
dataset = "scannetpp"
config = args.config
cropformer_path = args.cropformer_path
root = f'data/{config}/processed'
image_path_pattern = 'color/*.jpg' # stride = 10
# gt = 'data/scannet_dust3r_posed/gt'
t0 = time.time()
# seq_name_list = get_seq_name_list(config)
seq_name_list = os.listdir(root)
# val_path = Path("../") / "OKNO/data/arkitscenes/arkitscenes_offline_infos_val.pkl"
# out_dir = Path("data/arkit_vggt/processed")
# with open(val_path, "rb") as f:
# data_list = pickle.load(f)['data_list']
# val_scenes = [scene["lidar_points"]["lidar_path"] for scene in data_list][:2500]
# def extract_name(item):
# return item.split("_")[0]
# val_scenes = set([extract_name(scene) for scene in val_scenes])
# seq_name_list = list(seq_name_list.intersection(val_scenes))
# seq_name_list = ["scene0011_00", "scene0011_01", "scene0015_00", "scene0019_00", "scene0019_01", "scene0025_00", "scene0025_01", "scene0025_02", "scene0030_00", "scene0030_01", "scene0030_02", "scene0046_00", "scene0046_01", "scene0046_02", "scene0050_00", "scene0050_01", "scene0050_02", "scene0063_00", "scene0064_00", "scene0064_01", "scene0077_00", "scene0077_01", "scene0081_00", "scene0081_01", "scene0081_02", "scene0084_00", "scene0084_01", "scene0084_02", "scene0086_00", "scene0086_01", "scene0086_02", "scene0088_00", "scene0088_01", "scene0088_02", "scene0088_03", "scene0095_00", "scene0095_01", "scene0100_00", "scene0100_01", "scene0100_02", "scene0131_00", "scene0131_01", "scene0131_02", "scene0139_00", "scene0144_00", "scene0144_01", "scene0146_00", "scene0146_01", "scene0146_02", "scene0149_00", "scene0153_00", "scene0153_01", "scene0164_00", "scene0164_01", "scene0164_02", "scene0164_03", "scene0169_00", "scene0169_01", "scene0187_00", "scene0187_01", "scene0193_00", "scene0193_01", "scene0196_00", "scene0203_00", "scene0203_01", "scene0203_02", "scene0207_00", "scene0207_01", "scene0207_02", "scene0208_00", "scene0217_00", "scene0221_00", "scene0221_01", "scene0222_00", "scene0222_01", "scene0231_00", "scene0231_01", "scene0231_02", "scene0246_00", "scene0249_00", "scene0251_00", "scene0256_00", "scene0256_01", "scene0256_02", "scene0257_00", "scene0277_00", "scene0277_01", "scene0277_02", "scene0278_00", "scene0278_01", "scene0300_00", "scene0300_01", "scene0304_00", "scene0307_00", "scene0307_01", "scene0307_02", "scene0314_00", "scene0316_00", "scene0328_00", "scene0329_00", "scene0329_01", "scene0329_02", "scene0334_00", "scene0334_01", "scene0334_02", "scene0338_00", "scene0338_01", "scene0338_02", "scene0342_00", "scene0343_00", "scene0351_00", "scene0351_01", "scene0353_00", "scene0353_01", "scene0353_02", "scene0354_00", "scene0355_00", "scene0355_01", "scene0356_00", "scene0356_01", "scene0356_02", "scene0357_00", "scene0357_01", "scene0377_00", "scene0377_01", "scene0377_02", "scene0378_00", "scene0378_01", "scene0378_02", "scene0382_00", "scene0382_01", "scene0389_00", "scene0406_00", "scene0406_01", "scene0406_02", "scene0412_00", "scene0412_01", "scene0414_00", "scene0423_00", "scene0423_01", "scene0423_02", "scene0426_00", "scene0426_01", "scene0426_02", "scene0426_03", "scene0427_00", "scene0430_00", "scene0430_01", "scene0432_00", "scene0432_01", "scene0435_00", "scene0435_01", "scene0435_02", "scene0435_03", "scene0441_00", "scene0458_00", "scene0458_01", "scene0461_00", "scene0462_00", "scene0474_00", "scene0474_01", "scene0474_02", "scene0474_03", "scene0474_04", "scene0474_05", "scene0488_00", "scene0488_01", "scene0490_00", "scene0494_00", "scene0496_00", "scene0500_00", "scene0500_01", "scene0518_00", "scene0527_00", "scene0535_00", "scene0549_00", "scene0549_01", "scene0550_00", "scene0552_00", "scene0552_01", "scene0553_00", "scene0553_01", "scene0553_02", "scene0558_00", "scene0558_01", "scene0558_02", "scene0559_00", "scene0559_01", "scene0559_02", "scene0565_00", "scene0568_00", "scene0568_01", "scene0568_02", "scene0574_00", "scene0574_01", "scene0574_02", "scene0575_00", "scene0575_01", "scene0575_02", "scene0578_00", "scene0578_01", "scene0578_02", "scene0580_00", "scene0580_01", "scene0583_00", "scene0583_01", "scene0583_02", "scene0591_00", "scene0591_01", "scene0591_02", "scene0593_00", "scene0593_01", "scene0595_00", "scene0598_00", "scene0598_01", "scene0598_02", "scene0599_00", "scene0599_01", "scene0599_02", "scene0606_00", "scene0606_01", "scene0606_02", "scene0607_00", "scene0607_01", "scene0608_00", "scene0608_01", "scene0608_02", "scene0609_00", "scene0609_01", "scene0609_02", "scene0609_03", "scene0616_00", "scene0616_01", "scene0618_00", "scene0621_00", "scene0629_00", "scene0629_01", "scene0629_02", "scene0633_00", "scene0633_01", "scene0643_00", "scene0644_00", "scene0645_00", "scene0645_01", "scene0645_02", "scene0647_00", "scene0647_01", "scene0648_00", "scene0648_01", "scene0651_00", "scene0651_01", "scene0651_02", "scene0652_00", "scene0653_00", "scene0653_01", "scene0655_00", "scene0655_01", "scene0655_02", "scene0658_00", "scene0660_00", "scene0663_00", "scene0663_01", "scene0663_02", "scene0664_00", "scene0664_01", "scene0664_02", "scene0665_00", "scene0665_01", "scene0670_00", "scene0670_01", "scene0671_00", "scene0671_01", "scene0678_00", "scene0678_01", "scene0678_02", "scene0684_00", "scene0684_01", "scene0685_00", "scene0685_01", "scene0685_02", "scene0686_00", "scene0686_01", "scene0686_02", "scene0689_00", "scene0690_00", "scene0690_01", "scene0693_00", "scene0693_01", "scene0693_02", "scene0695_00", "scene0695_01", "scene0695_02", "scene0695_03", "scene0696_00", "scene0696_01", "scene0696_02", "scene0697_00", "scene0697_01", "scene0697_02", "scene0697_03", "scene0699_00", "scene0700_00", "scene0700_01", "scene0700_02", "scene0701_00", "scene0701_01", "scene0701_02", "scene0702_00", "scene0702_01", "scene0702_02", "scene0704_00", "scene0704_01"]
# print('There are %d scenes' % len(seq_name_list))
# Step 1: use Cropformer to get 2D instance masks for all sequences.
# parallel_compute(f'python third_party/detectron2/projects/CropFormer/demo_cropformer/mask_predict.py --config-file third_party/detectron2/projects/CropFormer/configs/entityv2/entity_segmentation/mask2former_hornet_3x.yaml --root {root} --image_path_pattern {image_path_pattern} --dataset {args.dataset} --seq_name_list %s --opts MODEL.WEIGHTS {cropformer_path}', 'predict mask', 'cuda', CUDA_LIST, seq_name_list)
t1 = time.time()
print('finish predict mask, cropformer time: ', t1 - t0, 'sec')
print('avg per scene time: ', (t1 - t0) / len(seq_name_list), 'sec')
# # Step 2: Mask clustering using our proposed method.
parallel_compute(f'python main.py --config {config} --seq_name_list %s', 'mask clustering', 'cuda', CUDA_LIST, seq_name_list)
t2 = time.time()
print('finish mask clustering, mask clustering time: ', t2 - t1)
print('avg per scene time: ', (t2 - t1) / len(seq_name_list))
# Step 3: Evaluate the class-agnostic results.
# os.system(f'python -m evaluation.evaluate --pred_path data/prediction/{config}_class_agnostic --gt_path {gt} --dataset {dataset} --no_class')
# Step 4: Get the open-vocabulary semantic features for each 2D masks.
parallel_compute(f'python -m semantics.get_open-voc_features --config {config} --seq_name_list %s', 'get open-vocabulary semantic features using CLIP', 'cuda', CUDA_LIST, seq_name_list)
# Step 5: Get the text CLIP features for each label.
# get_label_text_feature(CUDA_LIST[0])
# Step 6: Get labels for each 3D instances.
parallel_compute(f'python -m semantics.open-voc_query --config {config}', 'get text labels', 'cpu', CUDA_LIST, seq_name_list)
t3 = time.time()
print('finish get text labels, get text labels time: ', t3 - t2)
print('avg per scene time: ', (t3 - t2) / len(seq_name_list))
# Step 7: Evaluate the class-aware results.
# os.system(f'python -m evaluation.evaluate --pred_path data/prediction/{config} --gt_path {gt} --dataset {dataset}')
print('total time', (time.time() - t0)//60, 'min')
print('Average time', (time.time() - t0) / len(seq_name_list), 'sec')
if __name__ == '__main__':
args = get_args()
main(args)
|