Spaces:
Running
Running
File size: 2,094 Bytes
4c88343 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import argparse
from collections import defaultdict
from pathlib import Path
import numpy as np
from tqdm import tqdm
from . import logger
from .utils.read_write_model import read_model
def main(model, output, num_matched):
logger.info("Reading the COLMAP model...")
cameras, images, points3D = read_model(model)
logger.info("Extracting image pairs from covisibility info...")
pairs = []
for image_id, image in tqdm(images.items()):
matched = image.point3D_ids != -1
points3D_covis = image.point3D_ids[matched]
covis = defaultdict(int)
for point_id in points3D_covis:
for image_covis_id in points3D[point_id].image_ids:
if image_covis_id != image_id:
covis[image_covis_id] += 1
if len(covis) == 0:
logger.info(f"Image {image_id} does not have any covisibility.")
continue
covis_ids = np.array(list(covis.keys()))
covis_num = np.array([covis[i] for i in covis_ids])
if len(covis_ids) <= num_matched:
top_covis_ids = covis_ids[np.argsort(-covis_num)]
else:
# get covisible image ids with top k number of common matches
ind_top = np.argpartition(covis_num, -num_matched)
ind_top = ind_top[-num_matched:] # unsorted top k
ind_top = ind_top[np.argsort(-covis_num[ind_top])]
top_covis_ids = [covis_ids[i] for i in ind_top]
assert covis_num[ind_top[0]] == np.max(covis_num)
for i in top_covis_ids:
pair = (image.name, images[i].name)
pairs.append(pair)
logger.info(f"Found {len(pairs)} pairs.")
with open(output, "w") as f:
f.write("\n".join(" ".join([i, j]) for i, j in pairs))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, type=Path)
parser.add_argument("--output", required=True, type=Path)
parser.add_argument("--num_matched", required=True, type=int)
args = parser.parse_args()
main(**args.__dict__)
|