Spaces:
Running
Running
File size: 2,184 Bytes
4c88343 57c1094 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 62 63 64 65 66 67 |
import argparse
import collections.abc as collections
from pathlib import Path
from typing import List, Optional, Union
from . import logger
from .utils.io import list_h5_names
from .utils.parsers import parse_image_lists
def main(
output: Path,
image_list: Optional[Union[Path, List[str]]] = None,
features: Optional[Path] = None,
ref_list: Optional[Union[Path, List[str]]] = None,
ref_features: Optional[Path] = None,
):
if image_list is not None:
if isinstance(image_list, (str, Path)):
names_q = parse_image_lists(image_list)
elif isinstance(image_list, collections.Iterable):
names_q = list(image_list)
else:
raise ValueError(f"Unknown type for image list: {image_list}")
elif features is not None:
names_q = list_h5_names(features)
else:
raise ValueError("Provide either a list of images or a feature file.")
self_matching = False
if ref_list is not None:
if isinstance(ref_list, (str, Path)):
names_ref = parse_image_lists(ref_list)
elif isinstance(image_list, collections.Iterable):
names_ref = list(ref_list)
else:
raise ValueError(
f"Unknown type for reference image list: {ref_list}"
)
elif ref_features is not None:
names_ref = list_h5_names(ref_features)
else:
self_matching = True
names_ref = names_q
pairs = []
for i, n1 in enumerate(names_q):
for j, n2 in enumerate(names_ref):
if self_matching and j <= i:
continue
pairs.append((n1, n2))
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("--output", required=True, type=Path)
parser.add_argument("--image_list", type=Path)
parser.add_argument("--features", type=Path)
parser.add_argument("--ref_list", type=Path)
parser.add_argument("--ref_features", type=Path)
args = parser.parse_args()
main(**args.__dict__)
|