Function process_scenes

benchmark/reconstruction/evaluation/utils.py:831–929

(
    args: argparse.Namespace,
    scene_infos: list[SceneInfo],
    prepare_scene: Callable[[SceneInfo], None],
    position_accuracy_gt: float,
)

def process_scenes(
    args: argparse.Namespace,
    scene_infos: list[SceneInfo],
    prepare_scene: Callable[[SceneInfo], None],
    position_accuracy_gt: float,
) -> MetricsByCatByScene:
    error_thresholds = get_error_thresholds(args)

    gpu_index = _parse_gpu_index(args)
    scene_gpu_pairs = [
        (scene_info, str(gpu_index[i % len(gpu_index)]))
        for i, scene_info in enumerate(scene_infos)
    ]

    num_parallel_scenes = min(args.num_parallel_scenes, len(scene_infos))
    num_threads_per_scene = max(1, args.num_threads // num_parallel_scenes)

    manager = None
    progress_status = None
    monitor_thread = None
    stop_event = threading.Event()
    if args.progress:
        manager = multiprocessing.Manager()
        progress_status = manager.dict()
        monitor_thread = threading.Thread(
            target=_run_progress_monitor,
            args=(progress_status, len(scene_infos), stop_event),
            daemon=True,
        )
        monitor_thread.start()

    try:
        p = multiprocessing.Pool(
            processes=num_parallel_scenes, initializer=_init_pool_worker
        )
        try:
            results = list(
                p.imap_unordered(
                    functools.partial(
                        _process_scene_with_gpu,
                        args=args,
                        prepare_scene=prepare_scene,
                        position_accuracy_gt=position_accuracy_gt,
                        num_threads=num_threads_per_scene,
                        progress_status=progress_status,
                    ),
                    scene_gpu_pairs,
                    chunksize=1,
                )
            )
        except KeyboardInterrupt:
            pycolmap.logging.warning(
                "Interrupted, terminating workers and child processes..."
            )
            p.terminate()
            raise
        except BaseException:
            p.terminate()