Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Pipeline for custom data #315

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Pipeline for custom data #315

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,5 @@ __pycache__
.ipynb_checkpoints
outputs/
datasets/*
query_processing_data/
!datasets/sacre_coeur/
17 changes: 11 additions & 6 deletions hloc/localize_sfm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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...')
Expand Down Expand Up @@ -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']
Expand All @@ -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:
Expand Down
29 changes: 29 additions & 0 deletions hloc/pairs_from_retrieval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,35 @@ def pairs_from_score_matrix(scores: torch.Tensor,
return pairs


def save_global_candidates_for_query(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this function needed? why isn't pairs_from_retrieval.main sufficient?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right, pairs_from_retreival.main is sufficient. I will rewrite my example and push once I get the chance.

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):
Expand Down
272 changes: 272 additions & 0 deletions pipeline_custom.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}