Spaces:
Running
Running
File size: 2,477 Bytes
499e141 |
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 |
import argparse
import collections.abc as collections
import os
from pathlib import Path
from typing import Optional, Union, List
from . import logger
from .utils.parsers import parse_image_lists
from .utils.io import list_h5_names
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)):
if image_list.is_dir():
names_q = [x for x in os.listdir(str(image_list)) if x.endswith('.jpg') or x.endswith('.png')]
names_q.sort()
else:
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
# if j - i > 5:
# 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))
return 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__)
|