# -*- coding: UTF-8 -*- from pathlib import Path from typing import Tuple import cv2 import h5py import numpy as np from .parsers import names_to_pair, names_to_pair_old def read_image(path, grayscale=False): if grayscale: mode = cv2.IMREAD_GRAYSCALE else: mode = cv2.IMREAD_COLOR image = cv2.imread(str(path), mode) if image is None: raise ValueError(f"Cannot read image {path}.") if not grayscale and len(image.shape) == 3: image = image[:, :, ::-1] # BGR to RGB return image def list_h5_names(path): names = [] with h5py.File(str(path), "r", libver="latest") as fd: def visit_fn(_, obj): if isinstance(obj, h5py.Dataset): names.append(obj.parent.name.strip("/")) fd.visititems(visit_fn) return list(set(names)) def get_keypoints( path: Path, name: str, return_uncertainty: bool = False ) -> np.ndarray: with h5py.File(str(path), "r", libver="latest") as hfile: dset = hfile[name]["keypoints"] p = dset.__array__() uncertainty = dset.attrs.get("uncertainty") if return_uncertainty: return p, uncertainty return p def find_pair(hfile: h5py.File, name0: str, name1: str): pair = names_to_pair(name0, name1) if pair in hfile: return pair, False pair = names_to_pair(name1, name0) if pair in hfile: return pair, True # older, less efficient format pair = names_to_pair_old(name0, name1) if pair in hfile: return pair, False pair = names_to_pair_old(name1, name0) if pair in hfile: return pair, True raise ValueError( f"Could not find pair {(name0, name1)}... " "Maybe you matched with a different list of pairs? " ) def get_matches(path: Path, name0: str, name1: str) -> Tuple[np.ndarray]: with h5py.File(str(path), "r", libver="latest") as hfile: pair, reverse = find_pair(hfile, name0, name1) matches = hfile[pair]["matches0"].__array__() scores = hfile[pair]["matching_scores0"].__array__() idx = np.where(matches != -1)[0] matches = np.stack([idx, matches[idx]], -1) if reverse: matches = np.flip(matches, -1) scores = scores[idx] return matches, scores