From b4620beb0800a8cd632d032cb1eeaf6c254da3cd Mon Sep 17 00:00:00 2001 From: Sagar Bharadwaj Date: Fri, 6 Oct 2023 19:54:35 -0400 Subject: [PATCH] Custom pipeline --- .gitignore | 1 + hloc/localize_sfm.py | 17 ++- hloc/pairs_from_retrieval.py | 29 ++++ pipeline_custom.ipynb | 272 +++++++++++++++++++++++++++++++++++ 4 files changed, 313 insertions(+), 6 deletions(-) create mode 100644 pipeline_custom.ipynb diff --git a/.gitignore b/.gitignore index 64707625..e0b85381 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ __pycache__ .ipynb_checkpoints outputs/ datasets/* +query_processing_data/ !datasets/sacre_coeur/ diff --git a/hloc/localize_sfm.py b/hloc/localize_sfm.py index ba68d0e5..c79d2399 100644 --- a/hloc/localize_sfm.py +++ b/hloc/localize_sfm.py @@ -71,9 +71,13 @@ def pose_from_cluster( db_ids: List[int], features_path: Path, matches_path: Path, + features_q_path: Path = None, **kwargs): - kpq = get_keypoints(features_path, qname) + if features_q_path is None: + kpq = get_keypoints(features_path, qname) + else: + kpq = get_keypoints(features_q_path, qname) kpq += 0.5 # COLMAP coordinates kp_idx_to_3D = defaultdict(list) @@ -122,13 +126,13 @@ def pose_from_cluster( } return ret, log - def main(reference_sfm: Union[Path, pycolmap.Reconstruction], - queries: Path, + queries: Union[Path, list], retrieval: Path, features: Path, matches: Path, results: Path, + features_q: Path = None, ransac_thresh: int = 12, covisibility_clustering: bool = False, prepend_camera_name: bool = False, @@ -138,7 +142,8 @@ def main(reference_sfm: Union[Path, pycolmap.Reconstruction], assert features.exists(), features assert matches.exists(), matches - queries = parse_image_lists(queries, with_intrinsics=True) + if isinstance(queries, Path): + queries = parse_image_lists(queries, with_intrinsics=True) retrieval_dict = parse_retrieval(retrieval) logger.info('Reading the 3D model...') @@ -178,7 +183,7 @@ def main(reference_sfm: Union[Path, pycolmap.Reconstruction], logs_clusters = [] for i, cluster_ids in enumerate(clusters): ret, log = pose_from_cluster( - localizer, qname, qcam, cluster_ids, features, matches) + localizer, qname, qcam, cluster_ids, features, matches, features_q = features_q) if ret['success'] and ret['num_inliers'] > best_inliers: best_cluster = i best_inliers = ret['num_inliers'] @@ -194,7 +199,7 @@ def main(reference_sfm: Union[Path, pycolmap.Reconstruction], } else: ret, log = pose_from_cluster( - localizer, qname, qcam, db_ids, features, matches) + localizer, qname, qcam, db_ids, features, matches, features_q = features_q) if ret['success']: poses[qname] = (ret['qvec'], ret['tvec']) else: diff --git a/hloc/pairs_from_retrieval.py b/hloc/pairs_from_retrieval.py index 4329547a..88de1ef1 100644 --- a/hloc/pairs_from_retrieval.py +++ b/hloc/pairs_from_retrieval.py @@ -67,6 +67,35 @@ def pairs_from_score_matrix(scores: torch.Tensor, return pairs +def save_global_candidates_for_query( + db_descriptors, + query_descriptor, + query_image_names, + num_matched, + output_file_path + ): + """ + For each image in `query_image_names`, get a list of `num_matched` images from db that are closest to the + query image + """ + db_image_names = list_h5_names(db_descriptors) + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + db_desc = get_descriptors(db_image_names, db_descriptors) + query_desc = get_descriptors(query_image_names, query_descriptor) + sim = torch.einsum('id,jd->ij', query_desc.to(device), db_desc.to(device)) + + invalids = np.array(query_image_names)[:, None] == np.array(db_image_names)[None] + pairs = pairs_from_score_matrix(sim, invalids, num_matched, min_score=0) + pairs = [(query_image_names[i], db_image_names[j]) for i, j in pairs] + + logger.info(f'Found {len(pairs)} pairs.') + with open(output_file_path, 'w') as f: + f.write('\n'.join(' '.join([i, j]) for i, j in pairs)) + + if len(query_image_names) == 1: # If only one image is passed in, eeturn nearest candidates + return [p[1] for p in pairs] + def main(descriptors, output, num_matched, query_prefix=None, query_list=None, db_prefix=None, db_list=None, db_model=None, db_descriptors=None): diff --git a/pipeline_custom.ipynb b/pipeline_custom.ipynb new file mode 100644 index 00000000..72a016df --- /dev/null +++ b/pipeline_custom.ipynb @@ -0,0 +1,272 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "1aa2d1ce-7d86-4375-a491-540b56a1df90", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pycolmap\n", + "\n", + "from hloc import extract_features, pairs_from_covisibility, match_features, triangulation, pairs_from_retrieval, localize_sfm, visualization\n", + "from hloc.utils import viz_3d, io\n", + "from hloc.localize_sfm import QueryLocalizer, pose_from_cluster" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66d8ac9b-a21f-43df-bd2b-458fba523b5a", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "dataset_name = 'LabFront'\n", + "\n", + "dataset = Path(f'datasets/{dataset_name}/')\n", + "image_dir = dataset / 'images'\n", + "colmap_model_path = dataset / 'colmap/sparse/0'\n", + "\n", + "output_dir = dataset / 'hloc_data/'\n", + "# Local features and global descriptor filenames hardocded in feature confs\n", + "sfm_pairs_path = output_dir / 'sfm-pairs-covis20.txt' # Pairs used for SfM reconstruction\n", + "sfm_reconstruction_path = output_dir / 'sfm_reconstruction' # Path to reconstructed SfM\n", + "\n", + "# Query images\n", + "queries_img_dir = dataset / 'query'\n", + "\n", + "# Query data\n", + "query_processing_data_dir = Path(f'query_processing_data/{dataset_name}')\n", + "query_global_matches_path = query_processing_data_dir / 'global_match_pairs.txt'\n", + "query_local_match_path = query_processing_data_dir / 'local_match_data.h5'\n", + "query_results = query_processing_data_dir / 'query_results.txt'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aff8a17a-e45f-4a5a-9fc2-4d96e34edb57", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Feature extraction\n", + "## Extract local features in each data set image using Superpoint\n", + "local_feature_conf = extract_features.confs['superpoint_aachen']\n", + "local_features_path = extract_features.main(\n", + " conf = local_feature_conf,\n", + " image_dir = image_dir,\n", + " export_dir = output_dir\n", + ")\n", + "\n", + "## Extract global descriptors from each image using NetVLad\n", + "global_descriptor_conf = extract_features.confs['netvlad']\n", + "global_descriptors_path = extract_features.main(\n", + " conf = global_descriptor_conf,\n", + " image_dir = image_dir,\n", + " export_dir = output_dir\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0224548d-f5b9-448f-a997-5712f0207745", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Create SfM model using the local features just extracted\n", + "\n", + "## Note: There is already an SfM model created using Colmap available. However, that is created using the RootSIFT features.\n", + "## SfM model needs to be created using the new features.\n", + "\n", + "## Create matching pairs:\n", + "## Instead of creating image pairs by exhaustively searching through all possible pairs, we leverage the \n", + "## existing colmap model and form pairs by selecting the top 20 most covisibile neighbors for each image\n", + "pairs_from_covisibility.main(\n", + " model = colmap_model_path,\n", + " output = sfm_pairs_path,\n", + " num_matched = 20\n", + ")\n", + "\n", + "## Use the created pairs to match images and store the matching result in a match file\n", + "match_features_conf = match_features.confs['superglue']\n", + "sfm_matches_path = match_features.main(\n", + " conf = match_features_conf,\n", + " pairs = sfm_pairs_path,\n", + " features = local_feature_conf['output'], # This contains the file name where lcoal features are stored\n", + " export_dir = output_dir\n", + ")\n", + "\n", + "## Use the matches to reconstruct an SfM model\n", + "reconstruction = triangulation.main(\n", + " sfm_dir = sfm_reconstruction_path,\n", + " reference_model = colmap_model_path,\n", + " image_dir = image_dir,\n", + " pairs = sfm_pairs_path,\n", + " features = local_features_path,\n", + " matches = sfm_matches_path\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "696ab530-9124-4fd5-939a-95fdc7a694be", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Visualize the reconstruction\n", + "\n", + "fig = viz_3d.init_figure()\n", + "viz_3d.plot_reconstruction(fig, reconstruction, color='rgba(255,0,0,0.5)', name=\"mapping\", points_rgb=True)\n", + "fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b1123ae8-d2a7-4185-b689-082e2bc776dd", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Localize a new Image\n", + "\n", + "query_image_name = 'query_2.png'\n", + "\n", + "# Extarct local features and global descriptor for the new image\n", + "query_local_features_path = extract_features.main(\n", + " conf = local_feature_conf,\n", + " image_dir = queries_img_dir,\n", + " export_dir = query_processing_data_dir,\n", + " image_list = [query_image_name]\n", + ")\n", + "\n", + "query_global_descriptor_path = extract_features.main(\n", + " conf = global_descriptor_conf,\n", + " image_dir = queries_img_dir,\n", + " export_dir = query_processing_data_dir,\n", + " image_list = [query_image_name]\n", + ")\n", + "\n", + "## Use global descriptor matching to get candidate matches\n", + "nearest_candidate_images = pairs_from_retrieval.save_global_candidates_for_query(\n", + " db_descriptors = global_descriptors_path,\n", + " query_descriptor = query_global_descriptor_path,\n", + " query_image_names = [query_image_name],\n", + " num_matched = 10,\n", + " output_file_path = query_global_matches_path\n", + ")\n", + "\n", + "## Match the query image against the candidate pairs from above\n", + "match_features.match_from_paths(\n", + " conf = match_features_conf,\n", + " pairs_path = query_global_matches_path,\n", + " match_path = query_local_match_path,\n", + " feature_path_q = query_local_features_path,\n", + " feature_path_ref = local_features_path\n", + ")\n", + "\n", + "## Now we have global candidate and thier mathces. We use this, along with SfM reconstruction to localize the image.\n", + "\n", + "camera = pycolmap.infer_camera_from_image(queries_img_dir / query_image_name)\n", + "ref_ids = [reconstruction.find_image_with_name(r).image_id for r in nearest_candidate_images]\n", + "conf = {\n", + " 'estimation': {'ransac': {'max_error': 12}},\n", + " 'refinement': {'refine_focal_length': True, 'refine_extra_params': True},\n", + "}\n", + "localizer = QueryLocalizer(reconstruction, conf)\n", + "ret, log = pose_from_cluster(\n", + " localizer = localizer, \n", + " qname = query_image_name, \n", + " query_camera = camera, \n", + " db_ids = ref_ids, \n", + " features_path = local_features_path, \n", + " matches_path = query_local_match_path,\n", + " features_q_path = query_local_features_path\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca011477-2bf8-4d56-871d-0a53de76d952", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# visualization.visualize_loc_from_log(\n", + "# image_dir = queries_img_dir, \n", + "# query_name = query_image_name, \n", + "# loc = log, \n", + "# reconstruction = reconstruction, \n", + "# db_image_dir = image_dir\n", + "# )\n", + "\n", + "fig = viz_3d.init_figure()\n", + "viz_3d.plot_reconstruction(fig, reconstruction, color='rgba(255,0,0,0.5)', points_rgb=True)\n", + "pose = pycolmap.Image(tvec=ret['tvec'], qvec=ret['qvec'])\n", + "viz_3d.plot_camera_colmap(fig, pose, camera, color='rgba(0,255,0,0.5)', name=query_image_name, fill=True)\n", + "# visualize 2D-3D correspodences\n", + "inl_3d = np.array([reconstruction.points3D[pid].xyz for pid in np.array(log['points3D_ids'])[ret['inliers']]])\n", + "viz_3d.plot_points(fig, inl_3d, color=\"lime\", ps=1, name=query_image_name)\n", + "fig.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3b9ead3d-2098-4efa-b0e6-d388954e59bc", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3bead830-aedc-4893-b8fb-abc48841ab34", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.18" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}