demo.py 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. from argparse import ArgumentParser
  2. import cv2
  3. from algorithms.dense_optical_flow import dense_optical_flow
  4. from algorithms.lucas_kanade import lucas_kanade_method
  5. def main():
  6. parser = ArgumentParser()
  7. parser.add_argument(
  8. "--algorithm",
  9. choices=["farneback", "lucaskanade", "lucaskanade_dense", "rlof"],
  10. required=True,
  11. help="Optical flow algorithm to use",
  12. )
  13. parser.add_argument(
  14. "--video_path", default="videos/cat.mp4", help="Path to the video",
  15. )
  16. args = parser.parse_args()
  17. video_path = args.video_path
  18. if args.algorithm == "lucaskanade":
  19. lucas_kanade_method(video_path)
  20. elif args.algorithm == "lucaskanade_dense":
  21. method = cv2.optflow.calcOpticalFlowSparseToDense
  22. dense_optical_flow(method, video_path, to_gray=True)
  23. elif args.algorithm == "farneback":
  24. method = cv2.calcOpticalFlowFarneback
  25. params = [0.5, 3, 15, 3, 5, 1.2, 0] # Farneback's algorithm parameters
  26. dense_optical_flow(method, video_path, params, to_gray=True)
  27. elif args.algorithm == "rlof":
  28. method = cv2.optflow.calcOpticalFlowDenseRLOF
  29. dense_optical_flow(method, video_path)
  30. if __name__ == "__main__":
  31. main()