demo_toolbox.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. import argparse
  2. import os
  3. from pathlib import Path
  4. from toolbox import Toolbox
  5. from utils.argutils import print_args
  6. from utils.default_models import ensure_default_models
  7. if __name__ == '__main__':
  8. parser = argparse.ArgumentParser(
  9. description="Runs the toolbox.",
  10. formatter_class=argparse.ArgumentDefaultsHelpFormatter
  11. )
  12. parser.add_argument("-d", "--datasets_root", type=Path, help= \
  13. "Path to the directory containing your datasets. See toolbox/__init__.py for a list of "
  14. "supported datasets.", default=None)
  15. parser.add_argument("-m", "--models_dir", type=Path, default="saved_models",
  16. help="Directory containing all saved models")
  17. parser.add_argument("--cpu", action="store_true", help=\
  18. "If True, all inference will be done on CPU")
  19. parser.add_argument("--seed", type=int, default=None, help=\
  20. "Optional random number seed value to make toolbox deterministic.")
  21. args = parser.parse_args()
  22. arg_dict = vars(args)
  23. print_args(args, parser)
  24. # Hide GPUs from Pytorch to force CPU processing
  25. if arg_dict.pop("cpu"):
  26. os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
  27. # Remind the user to download pretrained models if needed
  28. ensure_default_models(args.models_dir)
  29. # Launch the toolbox
  30. Toolbox(**arg_dict)