Vincentqyw
add: files
9223079
raw
history blame
1.12 kB
import logging
import numpy as np
from hloc.utils.read_write_model import read_model, write_model
logger = logging.getLogger(__name__)
def create_reference_sfm(full_model, ref_model, blacklist=None, ext=".bin"):
"""Create a new COLMAP model with only training images."""
logger.info("Creating the reference model.")
ref_model.mkdir(exist_ok=True)
cameras, images, points3D = read_model(full_model, ext)
if blacklist is not None:
with open(blacklist, "r") as f:
blacklist = f.read().rstrip().split("\n")
images_ref = dict()
for id_, image in images.items():
if blacklist and image.name in blacklist:
continue
images_ref[id_] = image
points3D_ref = dict()
for id_, point3D in points3D.items():
ref_ids = [i for i in point3D.image_ids if i in images_ref]
if len(ref_ids) == 0:
continue
points3D_ref[id_] = point3D._replace(image_ids=np.array(ref_ids))
write_model(cameras, images_ref, points3D_ref, ref_model, ".bin")
logger.info(f"Kept {len(images_ref)} images out of {len(images)}.")