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)