From 57772a035d4c2cfacf900f3b913a6580fe37f6c4 Mon Sep 17 00:00:00 2001 From: tiantianxiangshang629 Date: Wed, 2 Apr 2025 08:24:14 -0700 Subject: [PATCH 1/5] Add binary tree partition library. Remove draft code. Add binayr tree partition result. Add OPENCV camera model. Untrack binary_tree_partition_visualization.py and update .gitignore Remove unused log. Fix format issue. Add docstring to binary tree partition. Move unittest under test folder. Remove incorrect unittest file. Use more informative function name. Add unittests for member functions. Fix import format. Remove tracked binary files and local image dataset Remove temp scripts and image assets from version control Remove points3D.txt from tracking --- .../binary_tree_partition.py | 185 ++++++++++++++++ gtsfm/runner/gtsfm_runner_base.py | 40 ++-- gtsfm/scene_optimizer.py | 20 +- gtsfm/utils/io.py | 11 +- tests/test_binary_tree_partition.py | 205 ++++++++++++++++++ 5 files changed, 436 insertions(+), 25 deletions(-) create mode 100644 gtsfm/graph_partitioner/binary_tree_partition.py create mode 100644 tests/test_binary_tree_partition.py diff --git a/gtsfm/graph_partitioner/binary_tree_partition.py b/gtsfm/graph_partitioner/binary_tree_partition.py new file mode 100644 index 000000000..faf6b0a12 --- /dev/null +++ b/gtsfm/graph_partitioner/binary_tree_partition.py @@ -0,0 +1,185 @@ +"""Implementation of a binary tree graph partitioner. + +This partitioner recursively partitions image pair graphs into a binary tree +structure up to a specified depth, using METIS-based ordering. Leaf nodes +represent explicit image keys and associated edge groupings. + +Authors: Shicong Ma +""" + +from typing import Dict, List, Tuple + +import gtsam +import networkx as nx +from gtsam import SymbolicFactorGraph + +import gtsfm.utils.logger as logger_utils +from gtsfm.graph_partitioner.graph_partitioner_base import GraphPartitionerBase + +logger = logger_utils.get_logger() + + +class BinaryTreeNode: + """Node class for a binary tree representing partitioned sets of image keys.""" + + def __init__(self, keys: List[int], depth: int): + """ + Initialize a BinaryTreeNode. + + Args: + keys: Image indices at this node (only populated at leaf level). + depth: Depth level in the binary tree. + """ + self.keys = keys # Only at leaves + self.left = None + self.right = None + self.depth = depth + + def is_leaf(self) -> bool: + """Check whether this node is a leaf node.""" + return self.left is None and self.right is None + + +class BinaryTreePartition(GraphPartitionerBase): + """Graph partitioner that uses a binary tree to recursively divide image pairs.""" + + def __init__(self, max_depth: int = 2): + """ + Initialize the BinaryTreePartition object. + + Args: + max_depth: Maximum depth of the binary tree; results in 2^depth partitions. + """ + super().__init__(process_name="BinaryTreePartition") + self.max_depth = max_depth + + def partition_image_pairs(self, image_pairs: List[Tuple[int, int]]) -> List[List[Tuple[int, int]]]: + """Partition image pairs into subgroups using a binary tree. + + Args: + image_pairs: List of image index pairs (i, j), where i < j. + + Returns: + A list of image pair subsets, one for each leaf in the binary tree. + """ + if not image_pairs: + logger.warning("No image pairs provided for partitioning.") + return [] + + symbol_graph, _, nx_graph = self._build_graphs(image_pairs) + ordering = gtsam.Ordering.MetisSymbolicFactorGraph(symbol_graph) + binary_tree_root_node = self._build_binary_partition(ordering) + + num_leaves = 2**self.max_depth + image_pairs_per_partition = [[] for _ in range(num_leaves)] + + partition_details = self._compute_leaf_partition_details(binary_tree_root_node, nx_graph) + + logger.info(f"BinaryTreePartition: partitioned into {len(partition_details)} leaf nodes.") + + for i in range(num_leaves): + edges_explicit = partition_details[i].get("edges_within_explicit", []) + edges_shared = partition_details[i].get("edges_with_shared", []) + image_pairs_per_partition[i] = edges_explicit + edges_shared + + for i, part in enumerate(partition_details): + explicit_keys = part.get("explicit_keys", []) + edges_within = part.get("edges_within_explicit", []) + edges_shared = part.get("edges_with_shared", []) + + logger.info( + f"Partition {i}:\n" + f" Explicit Image Keys that only exists within the current partition ({len(explicit_keys)}): {sorted(explicit_keys)}\n" + f" Internal Edges ({len(edges_within)}): {edges_within}\n" + f" Shared Edges ({len(edges_shared)})" + ) + + return image_pairs_per_partition + + def _build_graphs(self, image_pairs: List[Tuple[int, int]]) -> Tuple[SymbolicFactorGraph, List[int], nx.Graph]: + """Construct GTSAM and NetworkX graphs from image pairs. + + Args: + image_pairs: List of image index pairs. + + Returns: + A tuple of (SymbolicFactorGraph, list of keys, NetworkX graph). + """ + sfg = gtsam.SymbolicFactorGraph() + nxg = nx.Graph() + keys = set() + + for i, j in image_pairs: + key_i = gtsam.symbol("x", i) + key_j = gtsam.symbol("x", j) + keys.add(key_i) + keys.add(key_j) + + sfg.push_factor(key_i, key_j) + nxg.add_edge(key_i, key_j) + + return sfg, list(keys), nxg + + def _build_binary_partition(self, ordering: gtsam.Ordering) -> BinaryTreeNode: + """Build a binary tree of image keys based on a given ordering. + + Args: + ordering: GTSAM Ordering object created via METIS. + + Returns: + Root node of the binary tree. + """ + ordered_keys = [ordering.at(i) for i in range(ordering.size())] + + def split(keys: List[int], depth: int) -> BinaryTreeNode: + if depth == self.max_depth: + return BinaryTreeNode(keys, depth) + + mid = len(keys) // 2 + left_node = split(keys[:mid], depth + 1) + right_node = split(keys[mid:], depth + 1) + node = BinaryTreeNode([], depth) + node.left = left_node + node.right = right_node + return node + + return split(ordered_keys, 0) + + def _compute_leaf_partition_details( + self, + node: BinaryTreeNode, + nx_graph: nx.Graph, + ) -> List[Dict]: + """Recursively traverse the binary tree and return partition details per leaf. + + Args: + node: Current binary tree node being processed. + nx_graph: NetworkX graph built from image pairs. + + Returns: + A list of dictionaries containing partition details per leaf node. + """ + if node.is_leaf(): + explicit_keys = set(node.keys) + edges_within_explicit = [] + edges_with_shared = [] + + for u, v in nx_graph.edges(): + if u in explicit_keys and v in explicit_keys: + edges_within_explicit.append((gtsam.Symbol(u).index(), gtsam.Symbol(v).index())) + elif u in explicit_keys or v in explicit_keys: + edges_with_shared.append((gtsam.Symbol(u).index(), gtsam.Symbol(v).index())) + + return [ + { + "explicit_keys": [gtsam.Symbol(u).index() for u in explicit_keys], + "explicit_count": len(explicit_keys), + "edges_within_explicit": edges_within_explicit, + "edges_with_shared": edges_with_shared, + } + ] + + # Collect from both children + left_partitions = self._compute_leaf_partition_details(node.left, nx_graph) + right_partitions = self._compute_leaf_partition_details(node.right, nx_graph) + return left_partitions + right_partitions diff --git a/gtsfm/runner/gtsfm_runner_base.py b/gtsfm/runner/gtsfm_runner_base.py index 6e633a2fc..de0133f85 100644 --- a/gtsfm/runner/gtsfm_runner_base.py +++ b/gtsfm/runner/gtsfm_runner_base.py @@ -9,29 +9,29 @@ import dask import hydra import numpy as np - from dask import config as dask_config from dask.distributed import Client, LocalCluster, SSHCluster, performance_report -from gtsam import Rot3, Pose3, Unit3 +from gtsam import Pose3, Rot3, Unit3 from hydra.utils import instantiate from omegaconf import OmegaConf import gtsfm.evaluation.metrics_report as metrics_report -import gtsfm.utils.merging as merging_utils import gtsfm.utils.logger as logger_utils +import gtsfm.utils.merging as merging_utils import gtsfm.utils.metrics as metrics_utils import gtsfm.utils.viz as viz_utils -from gtsfm.common.gtsfm_data import GtsfmData from gtsfm import two_view_estimator +from gtsfm.common.gtsfm_data import GtsfmData from gtsfm.evaluation.metrics import GtsfmMetric, GtsfmMetricsGroup from gtsfm.frontend.correspondence_generator.image_correspondence_generator import ImageCorrespondenceGenerator +from gtsfm.graph_partitioner.binary_tree_partition import BinaryTreePartition +from gtsfm.graph_partitioner.graph_partitioner_base import GraphPartitionerBase +from gtsfm.graph_partitioner.single_partition import SinglePartition from gtsfm.loader.loader_base import LoaderBase from gtsfm.retriever.retriever_base import ImageMatchingRegime from gtsfm.scene_optimizer import SceneOptimizer from gtsfm.two_view_estimator import TWO_VIEW_OUTPUT, TwoViewEstimationReport, run_two_view_estimator_as_futures from gtsfm.ui.process_graph_generator import ProcessGraphGenerator -from gtsfm.graph_partitioner.graph_partitioner_base import GraphPartitionerBase -from gtsfm.graph_partitioner.single_partition import SinglePartition from gtsfm.utils.subgraph_utils import group_results_by_subgraph dask_config.set({"distributed.scheduler.worker-ttl": None}) @@ -55,6 +55,7 @@ def __init__(self, override_args=None) -> None: self.loader: LoaderBase = self.construct_loader() self.scene_optimizer: SceneOptimizer = self.construct_scene_optimizer() + self.graph_partitioner: GraphPartitionerBase = self.construct_graph_partitioner() def construct_argparser(self) -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=self.tag) @@ -158,7 +159,7 @@ def construct_argparser(self) -> argparse.ArgumentParser: "--graph_partitioner", type=str, default="single", - choices=["single", "other_partitioner_types"], + choices=["single", "binary_tree_eight_partitions", "other_partitioner_types"], help="Type of graph partitioner to use. Default is 'single' (SinglePartition).", ) return parser @@ -252,6 +253,12 @@ def construct_scene_optimizer(self) -> SceneOptimizer: logger.info("\n\nSceneOptimizer: " + str(scene_optimizer)) return scene_optimizer + def construct_graph_partitioner(self): + graph_partition_input = self.parsed_args.graph_partitioner + if graph_partition_input == "binary_tree_eight_partitions": + return BinaryTreePartition(max_depth=3) + return SinglePartition() + def setup_ssh_cluster_with_retries(self) -> SSHCluster: """Sets up SSH Cluster allowing multiple retries upon connection failures.""" workers = OmegaConf.load(os.path.join("gtsfm", "configs", self.parsed_args.cluster_config))["workers"] @@ -283,14 +290,10 @@ def setup_ssh_cluster_with_retries(self) -> SSHCluster: ) return cluster - def run(self, graph_partitioner: GraphPartitionerBase = None) -> GtsfmData: + def run(self) -> GtsfmData: """Run the SceneOptimizer.""" start_time = time.time() - # Create graph partitioner if not provided - if graph_partitioner is None: - graph_partitioner = SinglePartition() - # Create dask cluster. if self.parsed_args.cluster_config: cluster = self.setup_ssh_cluster_with_retries() @@ -322,7 +325,7 @@ def run(self, graph_partitioner: GraphPartitionerBase = None) -> GtsfmData: client=client, images=self.loader.get_all_images_as_futures(client), image_fnames=self.loader.image_filenames(), - plots_output_dir=self.scene_optimizer._plot_base_path, + plots_output_dir=self.scene_optimizer.create_plot_base_path(), ) retriever_metrics = self.scene_optimizer.image_pairs_generator._retriever.evaluate( @@ -394,9 +397,8 @@ def run(self, graph_partitioner: GraphPartitionerBase = None) -> GtsfmData: all_metrics_groups = [retriever_metrics, two_view_agg_metrics] # Partition image pairs - subgraphs = graph_partitioner.partition_image_pairs(image_pair_indices) + subgraphs = self.graph_partitioner.partition_image_pairs(image_pair_indices) logger.info(f"Partitioned into {len(subgraphs)} subgraphs") - # Group results by subgraph subgraph_two_view_results = group_results_by_subgraph(two_view_results_dict, subgraphs) @@ -407,9 +409,13 @@ def run(self, graph_partitioner: GraphPartitionerBase = None) -> GtsfmData: for idx, subgraph_result_dict in enumerate(subgraph_two_view_results): logger.info( - f"Creating computation graph for subgraph {idx+1}/{len(subgraph_two_view_results)}" - f"with {len(subgraph_result_dict)} image pairs" + f"Creating computation graph for subgraph {idx+1}/{len(subgraph_two_view_results)} with {len(subgraph_result_dict)} image pairs" ) + if len(subgraph_two_view_results) == 1: + # single partition + self.scene_optimizer.create_output_directories(None) + else: + self.scene_optimizer.create_output_directories(idx + 1) # Unzip the two-view results for this subgraph subgraph_i2Ri1_dict, subgraph_i2Ui1_dict, subgraph_v_corr_idxs_dict, _, subgraph_post_isp_reports = ( diff --git a/gtsfm/scene_optimizer.py b/gtsfm/scene_optimizer.py index ae8bcad6d..e1a878bbf 100644 --- a/gtsfm/scene_optimizer.py +++ b/gtsfm/scene_optimizer.py @@ -90,7 +90,6 @@ def __init__( self._pose_angular_error_thresh = pose_angular_error_thresh self.output_root = Path(output_root) self._output_worker = output_worker - self._create_output_directories() def __repr__(self) -> str: """Returns string representation of class.""" @@ -102,12 +101,21 @@ def __repr__(self) -> str: DenseMultiviewOptimizer: {self.dense_multiview_optimizer} """ - def _create_output_directories(self) -> None: + def create_plot_base_path(self): + """Create plot base path.""" + plot_base_path = self.output_root / "plots" + os.makedirs(plot_base_path, exist_ok=True) + return plot_base_path + + def create_output_directories(self, partition_index: Optional[int]) -> None: """Create various output directories for GTSFM results, metrics, and plots.""" - # base paths for storage - self._plot_base_path = self.output_root / "plots" - self._metrics_path = self.output_root / "result_metrics" - self._results_path = self.output_root / "results" + # Construct subfolder if partitioned + partition_folder = f"partition_{partition_index}" if partition_index is not None else "" + + # Base paths + self._plot_base_path = self.output_root / "plots" / partition_folder + self._metrics_path = self.output_root / "result_metrics" / partition_folder + self._results_path = self.output_root / "results" / partition_folder # plot paths self._plot_correspondence_path = self._plot_base_path / "correspondences" diff --git a/gtsfm/utils/io.py b/gtsfm/utils/io.py index 027d02bb7..ae00195cd 100644 --- a/gtsfm/utils/io.py +++ b/gtsfm/utils/io.py @@ -2,6 +2,7 @@ Authors: Ayush Baid, John Lambert """ + import glob import os import pickle @@ -15,7 +16,7 @@ import numpy as np import open3d import simplejson as json -from gtsam import Cal3Bundler, Point3, Pose3, Rot3, SfmTrack +from gtsam import Cal3Bundler, Cal3DS2, Point3, Pose3, Rot3, SfmTrack from PIL import Image as PILImage from PIL.ExifTags import GPSTAGS, TAGS @@ -223,10 +224,16 @@ def colmap2gtsfm( elif camera_model_name == "RADIAL": f, cx, cy, k1, k2 = cameras[img.camera_id].params[:5] fx = f + elif camera_model_name == "OPENCV": + fx, fy, cx, cy, k1, k2, p1, p2 = cameras[img.camera_id].params[:8] else: raise ValueError(f"Unsupported COLMAP camera type: {camera_model_name}") - intrinsics_gtsfm.append(Cal3Bundler(fx, k1, k2, cx, cy)) + if camera_model_name == "OPENCV": + intrinsics_gtsfm.append(Cal3DS2(fx, fy, 0.0, cx, cy, k1, k2, p1, p2)) + else: + intrinsics_gtsfm.append(Cal3Bundler(fx, k1, k2, cx, cy)) + image_id_to_idx[img.id] = idx img_h, img_w = cameras[img.camera_id].height, cameras[img.camera_id].width img_dims.append((img_h, img_w)) diff --git a/tests/test_binary_tree_partition.py b/tests/test_binary_tree_partition.py new file mode 100644 index 000000000..cb5f30f06 --- /dev/null +++ b/tests/test_binary_tree_partition.py @@ -0,0 +1,205 @@ +"""Unit tests for the BinaryTreePartition graph partitioner. + +This module tests the functionality of the BinaryTreePartition class, ensuring that: +- It correctly splits image pairs into leaf partitions. +- The resulting partitions contain valid and non-duplicated edges. +- The binary tree structure is valid. +- Edge cases like empty input are handled gracefully. + +Author: Shicong Ma +""" + +import unittest +from typing import List, Tuple + +import gtsam + +from gtsfm.graph_partitioner.binary_tree_partition import BinaryTreeNode, BinaryTreePartition + + +class TestBinaryTreePartition(unittest.TestCase): + """Unit tests for BinaryTreePartition.""" + + def setUp(self): + """Set up a simple 4x4 grid graph for testing.""" + self.rows, self.cols = 4, 4 + self.total_nodes = self.rows * self.cols + self.image_pairs = self._create_grid_edges(self.rows, self.cols) + + def _create_grid_edges(self, rows: int, cols: int) -> List[Tuple[int, int]]: + """Create a simple 2D grid of image pairs. + + Args: + rows: Number of rows in the grid. + cols: Number of columns in the grid. + + Returns: + List of edges between adjacent grid points. + """ + edges = [] + for r in range(rows): + for c in range(cols): + current_idx = r * cols + c + if c + 1 < cols: + edges.append((current_idx, current_idx + 1)) + if r + 1 < rows: + edges.append((current_idx, current_idx + cols)) + return edges + + def test_partition_leaf_count(self): + """Test that partitioner creates the correct number of leaf partitions.""" + partitioner = BinaryTreePartition(max_depth=2) + partitions = partitioner.partition_image_pairs(self.image_pairs) + expected_num_leaves = 2**partitioner.max_depth + self.assertEqual(len(partitions), expected_num_leaves) + + def test_no_duplicate_undirected_edges(self): + """Test that undirected edges are not duplicated (e.g., both (u,v) and (v,u)).""" + partitioner = BinaryTreePartition(max_depth=2) + partitions = partitioner.partition_image_pairs(self.image_pairs) + undirected_edges = set() + for partition in partitions: + for u, v in partition: + edge = (min(u, v), max(u, v)) + undirected_edges.add(edge) + self.assertEqual(len(undirected_edges), len(set(undirected_edges))) + + def test_edge_validity(self): + """Test that all edges are valid integer indices within the image grid.""" + partitioner = BinaryTreePartition(max_depth=2) + partitions = partitioner.partition_image_pairs(self.image_pairs) + for partition in partitions: + for u, v in partition: + self.assertIsInstance(u, int) + self.assertIsInstance(v, int) + self.assertGreaterEqual(u, 0) + self.assertLess(u, self.total_nodes) + self.assertGreaterEqual(v, 0) + self.assertLess(v, self.total_nodes) + + def test_non_empty_partitions(self): + """Test that at least one partition contains edges.""" + partitioner = BinaryTreePartition(max_depth=2) + partitions = partitioner.partition_image_pairs(self.image_pairs) + non_empty_count = sum(1 for p in partitions if len(p) > 0) + self.assertGreater(non_empty_count, 0) + + def test_empty_input(self): + """Test that empty image pair input returns an empty partition list.""" + partitioner = BinaryTreePartition(max_depth=2) + partitions = partitioner.partition_image_pairs([]) + self.assertEqual(partitions, []) + + def test_known_input_partition(self): + """Test partitioning of a simple known image pair set.""" + image_pairs = [(0, 1), (1, 2), (2, 3)] + partitioner = BinaryTreePartition(max_depth=1) + partitions = partitioner.partition_image_pairs(image_pairs) + self.assertEqual(len(partitions), 2) + + # All original edges should appear in at least one partition (either internal or shared) + flattened = set() + for p in partitions: + for u, v in p: + flattened.add((min(u, v), max(u, v))) + for u, v in image_pairs: + self.assertIn((min(u, v), max(u, v)), flattened) + + def test_binary_tree_structure(self): + """Test binary tree structure has correct number of leaves and leaf depths.""" + partitioner = BinaryTreePartition(max_depth=3) + + # Use synthetic ordering for test + ordering = type( + "FakeOrdering", + (), + {"size": lambda self: 8, "at": lambda self, idx: ord("a") + idx}, # returns int ASCII code + )() + + root = partitioner._build_binary_partition(ordering) + + leaf_nodes = [] + + def dfs(node): + if node.is_leaf(): + leaf_nodes.append(node) + if node.left: + dfs(node.left) + if node.right: + dfs(node.right) + + dfs(root) + self.assertEqual(len(leaf_nodes), 2**partitioner.max_depth) + self.assertTrue(all(n.depth == partitioner.max_depth for n in leaf_nodes)) + + def test_build_graphs(self): + """Test that _build_graphs constructs the correct symbolic and networkx graphs.""" + partitioner = BinaryTreePartition() + test_pairs = [(0, 1), (1, 2)] + sfg, keys, nxg = partitioner._build_graphs(test_pairs) + + # Check symbolic graph has correct number of factors + self.assertEqual(sfg.size(), len(test_pairs)) + + # Check nx graph has correct nodes and edges + self.assertEqual( + set(nxg.edges()), + {(gtsam.symbol("x", 0), gtsam.symbol("x", 1)), (gtsam.symbol("x", 1), gtsam.symbol("x", 2))}, + ) + self.assertEqual(len(nxg.nodes), 3) + self.assertEqual(len(keys), 3) + + def test_build_binary_partition(self): + """Test that binary tree is built correctly with specified depth and balanced splitting.""" + partitioner = BinaryTreePartition(max_depth=2) + + # Fake ordering: integers 0 through 7 + ordering = type("FakeOrdering", (), {"size": lambda self: 8, "at": lambda self, idx: idx})() + + root = partitioner._build_binary_partition(ordering) + + # Collect leaf nodes and check depth + leaf_nodes = [] + + def dfs(node): + if node.is_leaf(): + leaf_nodes.append(node) + if node.left: + dfs(node.left) + if node.right: + dfs(node.right) + + dfs(root) + self.assertEqual(len(leaf_nodes), 2**partitioner.max_depth) + self.assertTrue(all(n.depth == partitioner.max_depth for n in leaf_nodes)) + self.assertEqual(sum(len(n.keys) for n in leaf_nodes), 8) + + def test_compute_leaf_partition_details(self): + """Test that leaf partitions correctly report internal and shared edges.""" + partitioner = BinaryTreePartition(max_depth=1) + image_pairs = [(0, 1), (1, 2), (2, 3)] # Line graph + + # Build graph and partition tree + _, _, nxg = partitioner._build_graphs(image_pairs) + + # Manually create a binary tree with leaves split as [0,1] and [2,3] + left = BinaryTreeNode([gtsam.symbol("x", 0), gtsam.symbol("x", 1)], depth=1) + right = BinaryTreeNode([gtsam.symbol("x", 2), gtsam.symbol("x", 3)], depth=1) + root = BinaryTreeNode([], depth=0) + root.left = left + root.right = right + + details = partitioner._compute_leaf_partition_details(root, nxg) + self.assertEqual(len(details), 2) + + # Check that shared edge (1,2) is included in both + flattened_edges = set() + for d in details: + for u, v in d["edges_within_explicit"] + d["edges_with_shared"]: + flattened_edges.add((min(u, v), max(u, v))) + for u, v in image_pairs: + self.assertIn((min(u, v), max(u, v)), flattened_edges) + + +if __name__ == "__main__": + unittest.main() From 25232773c24b00dbf8b05fb8178211e7425f8073 Mon Sep 17 00:00:00 2001 From: tiantianxiangshang629 Date: Tue, 8 Jul 2025 22:30:37 -0700 Subject: [PATCH 2/5] Fix CI length issue. --- gtsfm/graph_partitioner/binary_tree_partition.py | 3 ++- gtsfm/runner/gtsfm_runner_base.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/gtsfm/graph_partitioner/binary_tree_partition.py b/gtsfm/graph_partitioner/binary_tree_partition.py index faf6b0a12..99a3550b7 100644 --- a/gtsfm/graph_partitioner/binary_tree_partition.py +++ b/gtsfm/graph_partitioner/binary_tree_partition.py @@ -89,7 +89,8 @@ def partition_image_pairs(self, image_pairs: List[Tuple[int, int]]) -> List[List logger.info( f"Partition {i}:\n" - f" Explicit Image Keys that only exists within the current partition ({len(explicit_keys)}): {sorted(explicit_keys)}\n" + f" Explicit Image Keys that only exist within the current partition " + f"({len(explicit_keys)}): {sorted(explicit_keys)}\n" f" Internal Edges ({len(edges_within)}): {edges_within}\n" f" Shared Edges ({len(edges_shared)})" ) diff --git a/gtsfm/runner/gtsfm_runner_base.py b/gtsfm/runner/gtsfm_runner_base.py index de0133f85..74aa67986 100644 --- a/gtsfm/runner/gtsfm_runner_base.py +++ b/gtsfm/runner/gtsfm_runner_base.py @@ -409,7 +409,8 @@ def run(self) -> GtsfmData: for idx, subgraph_result_dict in enumerate(subgraph_two_view_results): logger.info( - f"Creating computation graph for subgraph {idx+1}/{len(subgraph_two_view_results)} with {len(subgraph_result_dict)} image pairs" + f"Creating computation graph for subgraph {idx + 1}/{len(subgraph_two_view_results)} " + f"with { len(subgraph_result_dict)} image pairs" ) if len(subgraph_two_view_results) == 1: # single partition From c398e0db48cb5bbefa93db4fb5192ed7bbf5d422 Mon Sep 17 00:00:00 2001 From: tiantianxiangshang629 Date: Wed, 9 Jul 2025 00:37:16 -0700 Subject: [PATCH 3/5] Fix shared edge logic. --- .../binary_tree_partition.py | 38 ++++++++++++------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/gtsfm/graph_partitioner/binary_tree_partition.py b/gtsfm/graph_partitioner/binary_tree_partition.py index 99a3550b7..e63d31a4f 100644 --- a/gtsfm/graph_partitioner/binary_tree_partition.py +++ b/gtsfm/graph_partitioner/binary_tree_partition.py @@ -7,7 +7,7 @@ Authors: Shicong Ma """ -from typing import Dict, List, Tuple +from typing import Dict, List, Set, Tuple import gtsam import networkx as nx @@ -92,7 +92,7 @@ def partition_image_pairs(self, image_pairs: List[Tuple[int, int]]) -> List[List f" Explicit Image Keys that only exist within the current partition " f"({len(explicit_keys)}): {sorted(explicit_keys)}\n" f" Internal Edges ({len(edges_within)}): {edges_within}\n" - f" Shared Edges ({len(edges_shared)})" + f" Shared Edges ({len(edges_shared)}): {edges_shared}\n" ) return image_pairs_per_partition @@ -162,25 +162,35 @@ def _compute_leaf_partition_details( """ if node.is_leaf(): explicit_keys = set(node.keys) - edges_within_explicit = [] - edges_with_shared = [] - - for u, v in nx_graph.edges(): - if u in explicit_keys and v in explicit_keys: - edges_within_explicit.append((gtsam.Symbol(u).index(), gtsam.Symbol(v).index())) - elif u in explicit_keys or v in explicit_keys: - edges_with_shared.append((gtsam.Symbol(u).index(), gtsam.Symbol(v).index())) - return [ { "explicit_keys": [gtsam.Symbol(u).index() for u in explicit_keys], "explicit_count": len(explicit_keys), - "edges_within_explicit": edges_within_explicit, - "edges_with_shared": edges_with_shared, + "edges_within_explicit": [ + (gtsam.Symbol(u).index(), gtsam.Symbol(v).index()) + for u, v in nx_graph.edges() + if u in explicit_keys and v in explicit_keys + ], + "edges_with_shared": [], # placeholder } ] - # Collect from both children + # Recursively compute for children left_partitions = self._compute_leaf_partition_details(node.left, nx_graph) right_partitions = self._compute_leaf_partition_details(node.right, nx_graph) + + if node.left.is_leaf() and node.right.is_leaf(): + left_keys = set(node.left.keys) + right_keys = set(node.right.keys) + + shared_edges = [ + (gtsam.Symbol(u).index(), gtsam.Symbol(v).index()) + for u, v in nx_graph.edges() + if (u in left_keys and v in right_keys) or (u in right_keys and v in left_keys) + ] + + # Directly assign shared edges to the only two leaf partitions + left_partitions[0]["edges_with_shared"] = shared_edges + right_partitions[0]["edges_with_shared"] = shared_edges + return left_partitions + right_partitions From 1342e420c90328f2a9b21582188ddafd466df458 Mon Sep 17 00:00:00 2001 From: tiantianxiangshang629 Date: Wed, 9 Jul 2025 01:07:46 -0700 Subject: [PATCH 4/5] Add graph partition to scene optimizer. --- gtsfm/runner/gtsfm_runner_base.py | 10 ++-------- gtsfm/scene_optimizer.py | 4 ++++ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/gtsfm/runner/gtsfm_runner_base.py b/gtsfm/runner/gtsfm_runner_base.py index 74aa67986..ce102c16c 100644 --- a/gtsfm/runner/gtsfm_runner_base.py +++ b/gtsfm/runner/gtsfm_runner_base.py @@ -55,7 +55,7 @@ def __init__(self, override_args=None) -> None: self.loader: LoaderBase = self.construct_loader() self.scene_optimizer: SceneOptimizer = self.construct_scene_optimizer() - self.graph_partitioner: GraphPartitionerBase = self.construct_graph_partitioner() + self.graph_partitioner: GraphPartitionerBase = self.scene_optimizer.graph_partitioner def construct_argparser(self) -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=self.tag) @@ -159,7 +159,7 @@ def construct_argparser(self) -> argparse.ArgumentParser: "--graph_partitioner", type=str, default="single", - choices=["single", "binary_tree_eight_partitions", "other_partitioner_types"], + choices=["single", "other_partitioner_types"], help="Type of graph partitioner to use. Default is 'single' (SinglePartition).", ) return parser @@ -253,12 +253,6 @@ def construct_scene_optimizer(self) -> SceneOptimizer: logger.info("\n\nSceneOptimizer: " + str(scene_optimizer)) return scene_optimizer - def construct_graph_partitioner(self): - graph_partition_input = self.parsed_args.graph_partitioner - if graph_partition_input == "binary_tree_eight_partitions": - return BinaryTreePartition(max_depth=3) - return SinglePartition() - def setup_ssh_cluster_with_retries(self) -> SSHCluster: """Sets up SSH Cluster allowing multiple retries upon connection failures.""" workers = OmegaConf.load(os.path.join("gtsfm", "configs", self.parsed_args.cluster_config))["workers"] diff --git a/gtsfm/scene_optimizer.py b/gtsfm/scene_optimizer.py index e1a878bbf..cf9c68563 100644 --- a/gtsfm/scene_optimizer.py +++ b/gtsfm/scene_optimizer.py @@ -31,6 +31,8 @@ from gtsfm.common.pose_prior import PosePrior from gtsfm.densify.mvs_base import MVSBase from gtsfm.frontend.correspondence_generator.correspondence_generator_base import CorrespondenceGeneratorBase +from gtsfm.graph_partitioner.graph_partitioner_base import GraphPartitionerBase +from gtsfm.graph_partitioner.single_partition import SinglePartition from gtsfm.multi_view_optimizer import MultiViewOptimizer from gtsfm.retriever.image_pairs_generator import ImagePairsGenerator from gtsfm.retriever.retriever_base import ImageMatchingRegime @@ -75,6 +77,7 @@ def __init__( pose_angular_error_thresh: float = 3, # in degrees output_root: str = DEFAULT_OUTPUT_ROOT, output_worker: Optional[str] = None, + graph_partitioner: Optional[GraphPartitionerBase] = SinglePartition(), ) -> None: self.image_pairs_generator = image_pairs_generator self.correspondence_generator = correspondence_generator @@ -90,6 +93,7 @@ def __init__( self._pose_angular_error_thresh = pose_angular_error_thresh self.output_root = Path(output_root) self._output_worker = output_worker + self.graph_partitioner = graph_partitioner def __repr__(self) -> str: """Returns string representation of class.""" From 96d191c03aae8b5ed7f4dc5b9261fc61da511197 Mon Sep 17 00:00:00 2001 From: tiantianxiangshang629 Date: Wed, 9 Jul 2025 01:19:01 -0700 Subject: [PATCH 5/5] Fix CI format. --- gtsfm/graph_partitioner/binary_tree_partition.py | 2 +- gtsfm/runner/gtsfm_runner_base.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/gtsfm/graph_partitioner/binary_tree_partition.py b/gtsfm/graph_partitioner/binary_tree_partition.py index e63d31a4f..b42f08ba0 100644 --- a/gtsfm/graph_partitioner/binary_tree_partition.py +++ b/gtsfm/graph_partitioner/binary_tree_partition.py @@ -7,7 +7,7 @@ Authors: Shicong Ma """ -from typing import Dict, List, Set, Tuple +from typing import Dict, List, Tuple import gtsam import networkx as nx diff --git a/gtsfm/runner/gtsfm_runner_base.py b/gtsfm/runner/gtsfm_runner_base.py index ce102c16c..c3483fb33 100644 --- a/gtsfm/runner/gtsfm_runner_base.py +++ b/gtsfm/runner/gtsfm_runner_base.py @@ -24,9 +24,7 @@ from gtsfm.common.gtsfm_data import GtsfmData from gtsfm.evaluation.metrics import GtsfmMetric, GtsfmMetricsGroup from gtsfm.frontend.correspondence_generator.image_correspondence_generator import ImageCorrespondenceGenerator -from gtsfm.graph_partitioner.binary_tree_partition import BinaryTreePartition from gtsfm.graph_partitioner.graph_partitioner_base import GraphPartitionerBase -from gtsfm.graph_partitioner.single_partition import SinglePartition from gtsfm.loader.loader_base import LoaderBase from gtsfm.retriever.retriever_base import ImageMatchingRegime from gtsfm.scene_optimizer import SceneOptimizer