File size: 1,847 Bytes
6979012 | 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 | import argparse
from superpoint_pruning.evaluation import eval as eval_mod
from superpoint_pruning import export as export_mod
from superpoint_pruning.distillation import setup as setup_mod
from superpoint_pruning.distillation import lightning_trainer as train_mod
from superpoint_pruning.evaluation import benchmark_sp_tensorrt as benchmark_mod
from superpoint_pruning.evaluation import plot_keypoints as plot_keypoints_mod
def main():
parser = argparse.ArgumentParser(prog="superpoint-pruning")
subparsers = parser.add_subparsers(dest="command", required=True)
eval_parser = subparsers.add_parser("evaluate", help="Run SuperPoint evaluation")
eval_mod.add_parser_args(eval_parser)
eval_parser.set_defaults(func=eval_mod.main)
export_parser = subparsers.add_parser("export", help="Export SuperPoint to ONNX")
export_mod.add_parser_args(export_parser)
export_parser.set_defaults(func=export_mod.main)
setup_parser = subparsers.add_parser(
"setup", help="Setup the dataset for distillation and evaluation"
)
setup_mod.add_parser_args(setup_parser)
setup_parser.set_defaults(func=setup_mod.main)
train_parser = subparsers.add_parser("train", help="Train a pruned SuperPoint")
train_mod.add_parser_args(train_parser)
train_parser.set_defaults(func=train_mod.main)
benchmark_parser = subparsers.add_parser(
"benchmark", help="Benchmark a TensorRT SuperPoint engine"
)
benchmark_mod.add_parser_args(benchmark_parser)
benchmark_parser.set_defaults(func=benchmark_mod.main)
plot_parser = subparsers.add_parser(
"plot", help="Plot original and pruned SuperPoint keypoints"
)
plot_keypoints_mod.add_parser_args(plot_parser)
plot_parser.set_defaults(func=plot_keypoints_mod.main)
args = parser.parse_args()
args.func(args)
|