MaskClustering/run_scannet.py

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)