74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from .ccip import (
|
|
ccip_extract_feature,
|
|
ccip_batch_extract_features,
|
|
ccip_difference,
|
|
ccip_same,
|
|
)
|
|
except Exception:
|
|
sys.path.append(str(Path(__file__).resolve().parent.parent))
|
|
from ccip_lib.ccip import (
|
|
ccip_extract_feature,
|
|
ccip_batch_extract_features,
|
|
ccip_difference,
|
|
ccip_same,
|
|
)
|
|
|
|
|
|
def _parse_args():
|
|
p = argparse.ArgumentParser(prog='ccip-cli', description='Simple CCIP inference CLI (local model folder supported)')
|
|
sub = p.add_subparsers(dest='cmd')
|
|
|
|
ex = sub.add_parser('extract', help='Extract feature of a single image')
|
|
ex.add_argument('image')
|
|
ex.add_argument('--model-dir', '-m', help='Local model folder (containing model_feat.onnx etc.)')
|
|
ex.add_argument('--size', type=int, default=384)
|
|
ex.add_argument('--out', '-o', help='Save feature to .npy file')
|
|
|
|
diff = sub.add_parser('diff', help='Compute difference between two images')
|
|
diff.add_argument('image1')
|
|
diff.add_argument('image2')
|
|
diff.add_argument('--model-dir', '-m')
|
|
diff.add_argument('--size', type=int, default=384)
|
|
|
|
same = sub.add_parser('same', help='Decide whether two images are same character')
|
|
same.add_argument('image1')
|
|
same.add_argument('image2')
|
|
same.add_argument('--model-dir', '-m')
|
|
same.add_argument('--size', type=int, default=384)
|
|
same.add_argument('--threshold', type=float)
|
|
|
|
return p.parse_args()
|
|
|
|
|
|
def main(argv=None):
|
|
args = _parse_args() if argv is None else _parse_args()
|
|
|
|
model = args.model_dir if hasattr(args, 'model_dir') and args.model_dir else None
|
|
|
|
if args.cmd == 'extract':
|
|
feat = ccip_extract_feature(args.image, size=args.size, model=model or None)
|
|
print(feat.shape, feat.dtype)
|
|
if args.out:
|
|
import numpy as _np
|
|
_np.save(args.out, feat)
|
|
|
|
elif args.cmd == 'diff':
|
|
diff = ccip_difference(args.image1, args.image2, size=args.size, model=model or None)
|
|
print(diff)
|
|
|
|
elif args.cmd == 'same':
|
|
res = ccip_same(args.image1, args.image2, threshold=getattr(args, 'threshold', None), size=args.size, model=model or None)
|
|
print(res)
|
|
|
|
else:
|
|
print('No command. Use -h for help.')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|