Spaces:
Running
Running
File size: 2,282 Bytes
63f3cf2 |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
# -*- 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
|