Alleviate HDF5 bottleneck

hloc/extract_features.py

@@ -259,7 +259,7 @@ def main(conf: Dict,
                 if (dt == np.float32) and (dt != np.float16):
                     pred[k] = pred[k].astype(np.float16)
 
-        with h5py.File(str(feature_path), 'a') as fd:
+        with h5py.File(str(feature_path), 'a', libver='latest') as fd:
             try:
                 if name in fd:
                     del fd[name]

hloc/localize_inloc.py

@@ -121,8 +121,8 @@ def main(dataset_dir, retrieval, features, matches, results,
     retrieval_dict = parse_retrieval(retrieval)
     queries = list(retrieval_dict.keys())
 
-    feature_file = h5py.File(features, 'r')
-    match_file = h5py.File(matches, 'r')
+    feature_file = h5py.File(features, 'r', libver='latest')
+    match_file = h5py.File(matches, 'r', libver='latest')
 
     poses = {}
     logs = {

hloc/match_features.py

@@ -103,7 +103,7 @@ def find_unique_new_pairs(pairs_all: List[Tuple[str]], match_path: Path = None):
             pairs.add((i, j))
     pairs = list(pairs)
     if match_path is not None and match_path.exists():
-        with h5py.File(str(match_path), 'r') as fd:
+        with h5py.File(str(match_path), 'r', libver='latest') as fd:
             pairs_filtered = []
             for i, j in pairs:
                 if (names_to_pair(i, j) in fd or
@@ -149,13 +149,13 @@ def match_from_paths(conf: Dict,
 
     for (name0, name1) in tqdm(pairs, smoothing=.1):
         data = {}
-        with h5py.File(str(feature_path_q), 'r') as fd:
+        with h5py.File(str(feature_path_q), 'r', libver='latest') as fd:
             grp = fd[name0]
             for k, v in grp.items():
                 data[k+'0'] = torch.from_numpy(v.__array__()).float().to(device)
             # some matchers might expect an image but only use its size
             data['image0'] = torch.empty((1,)+tuple(grp['image_size'])[::-1])
-        with h5py.File(str(feature_paths_refs[name2ref[name1]]), 'r') as fd:
+        with h5py.File(str(feature_paths_refs[name2ref[name1]]), 'r', libver='latest') as fd:
             grp = fd[name1]
             for k, v in grp.items():
                 data[k+'1'] = torch.from_numpy(v.__array__()).float().to(device)
@@ -164,7 +164,7 @@ def match_from_paths(conf: Dict,
 
         pred = model(data)
         pair = names_to_pair(name0, name1)
-        with h5py.File(str(match_path), 'a') as fd:
+        with h5py.File(str(match_path), 'a', libver='latest') as fd:
             if pair in fd:
                 del fd[pair]
             grp = fd.create_group(pair)

hloc/pairs_from_retrieval.py

@@ -32,12 +32,12 @@ def parse_names(prefix, names, names_all):
 
 def get_descriptors(names, path, name2idx=None, key='global_descriptor'):
     if name2idx is None:
-        with h5py.File(str(path), 'r') as fd:
+        with h5py.File(str(path), 'r', libver='latest') as fd:
             desc = [fd[n][key].__array__() for n in names]
     else:
         desc = []
         for n in names:
-            with h5py.File(str(path[name2idx[n]]), 'r') as fd:
+            with h5py.File(str(path[name2idx[n]]), 'r', libver='latest') as fd:
                 desc.append(fd[n][key].__array__())
     return torch.from_numpy(np.stack(desc, 0)).float()
 

hloc/utils/io.py

@@ -22,7 +22,7 @@ def read_image(path, grayscale=False):
 
 def list_h5_names(path):
     names = []
-    with h5py.File(str(path), 'r') as fd:
+    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('/'))
@@ -32,7 +32,7 @@ def visit_fn(_, obj):
 
 def get_keypoints(path: Path, name: str,
                   return_uncertainty: bool = False) -> np.ndarray:
-    with h5py.File(str(path), 'r') as hfile:
+    with h5py.File(str(path), 'r', libver='latest') as hfile:
         dset = hfile[name]['keypoints']
         p = dset.__array__()
         uncertainty = dset.attrs.get('uncertainty')
@@ -61,7 +61,7 @@ def find_pair(hfile: h5py.File, name0: str, name1: str):
 
 
 def get_matches(path: Path, name0: str, name1: str) -> Tuple[np.ndarray]:
-    with h5py.File(str(path), 'r') as hfile:
+    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__()