File size: 7,582 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
import os
from tqdm import tqdm
import time
from utils.config import get_args
import argparse
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(config):
# Map config names to actual data directories
config_dir = config
if config in ['scannet18', 'scannet_dust3r_posed', 'scannet_dust3r_unposed',
'scannet_dust3r_posed_15', 'scannet_dust3r_posed_25', 'scannet_dust3r_posed_35',
'scannet_dust3r_unposed_15', 'scannet_dust3r_unposed_25', 'scannet_dust3r_unposed_35']:
if config == 'scannet18':
config_dir = 'scannet'
else:
config_dir = config
elif config in ['scannetpp_v2_dust3r_posed', 'scannetpp_v2_dust3r_unposed']:
config_dir = config
root = f'data/{config_dir}/processed'
if not os.path.exists(root):
raise FileNotFoundError(f"Directory not found: {root}")
seq_name_list = os.listdir(root)
return seq_name_list
def filter_completed_step1(config, seq_name_list):
"""Filter scenes that already have CropFormer masks"""
config_dir = config if config != 'scannet18' else 'scannet'
root = f'data/{config_dir}/processed'
filtered = []
for seq_name in seq_name_list:
mask_dir = os.path.join(root, seq_name, 'output/mask')
# Check if mask directory exists and has mask files
if not os.path.exists(mask_dir) or len(os.listdir(mask_dir)) == 0:
filtered.append(seq_name)
print(f'Step 1 (CropFormer): {len(filtered)}/{len(seq_name_list)} scenes need processing')
return filtered
def filter_completed_step2(config, seq_name_list):
"""Filter scenes that already have mask clustering results"""
config_dir = config if config != 'scannet18' else 'scannet'
root = f'data/{config_dir}/processed'
filtered = []
for seq_name in seq_name_list:
object_file = os.path.join(root, seq_name, 'output/object/object_dict.pkl')
if not os.path.exists(object_file):
filtered.append(seq_name)
print(f'Step 2 (Mask Clustering): {len(filtered)}/{len(seq_name_list)} scenes need processing')
return filtered
def filter_completed_step3(config, seq_name_list):
"""Filter scenes that already have CLIP features"""
config_dir = config if config != 'scannet18' else 'scannet'
root = f'data/{config_dir}/processed'
filtered = []
for seq_name in seq_name_list:
features_file = os.path.join(root, seq_name, 'output/features_clip.npy')
if not os.path.exists(features_file):
filtered.append(seq_name)
print(f'Step 3 (CLIP Features): {len(filtered)}/{len(seq_name_list)} scenes need processing')
return filtered
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 main(args):
config = args.config
cuda_list = args.cuda_list
cropformer_path = args.cropformer_path
dataset = args.dataset if args.dataset else config
# Map config names to actual data directories
config_dir = config
if config == 'scannet18':
config_dir = 'scannet'
root = f'data/{config_dir}/processed'
image_path_pattern = 'color/*.jpg'
t0 = time.time()
# Check if the processed directory exists
if not os.path.exists(root):
print(f'Processed directory not found: {root}')
print('Please run export scripts first!')
return
seq_name_list = get_seq_name_list(config)
print(f'There are {len(seq_name_list)} scenes exported and ready to process in {config}')
# Step 1: use Cropformer to get 2D instance masks for all sequences.
seq_list_step1 = filter_completed_step1(config, seq_name_list)
if len(seq_list_step1) > 0:
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 {dataset} --seq_name_list %s --opts MODEL.WEIGHTS {cropformer_path}', 'predict mask', 'cuda', cuda_list, seq_list_step1)
else:
print('Step 1: All scenes already have CropFormer masks, skipping...')
# Step 2: Mask clustering using our proposed method.
seq_list_step2 = filter_completed_step2(config, seq_name_list)
if len(seq_list_step2) > 0:
parallel_compute(f'python main.py --config {config} --seq_name_list %s', 'mask clustering', 'cuda', cuda_list, seq_list_step2)
else:
print('Step 2: All scenes already have mask clustering results, skipping...')
# Step 3: Get the open-vocabulary semantic features for each 2D masks.
seq_list_step3 = filter_completed_step3(config, seq_name_list)
if len(seq_list_step3) > 0:
parallel_compute(f'python -m semantics.get_open-voc_features --config {config} --dataset {config} --seq_name_list %s', 'get open-vocabulary semantic features using CLIP', 'cuda', cuda_list, seq_list_step3)
else:
print('Step 3: All scenes already have CLIP features, skipping...')
# Step 4: Get labels for each 3D instances.
parallel_compute(f'python -m semantics.open-voc_query --config {config} --dataset {config}', 'get text labels', 'cpu', cuda_list, seq_name_list)
print('total time', (time.time() - t0)//60, 'min')
print('Average time', (time.time() - t0) / len(seq_name_list), 'sec')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Run mask clustering on ScanNet datasets')
parser.add_argument('--config', type=str, required=True,
choices=['scannet18', 'scannet_dust3r_posed_15', 'scannet_dust3r_posed_25', 'scannet_dust3r_posed_35',
'scannet_dust3r_unposed_15', 'scannet_dust3r_unposed_25', 'scannet_dust3r_unposed_35',
'scannetpp_v2_dust3r_posed', 'scannetpp_v2_dust3r_unposed'],
help='Config name for the dataset')
parser.add_argument('--dataset', type=str, default=None,
help='Dataset name to pass to CropFormer (defaults to config name)')
parser.add_argument('--cuda_list', type=int, nargs='+', default=[0, 1, 2, 3],
help='List of CUDA device IDs to use (e.g., --cuda_list 0 1 2 3)')
parser.add_argument('--cropformer_path', type=str, default='Mask2Former_hornet_3x_576d0b.pth',
help='Path to CropFormer model weights')
args = parser.parse_args()
main(args)
|