FileUtils.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. # This file defines helper functions for processing files.
  2. from glob import glob
  3. import os
  4. import pickle
  5. expected_subfolders = sorted(["Motion", "Lapse", "Full"])
  6. def list_folders(path: str) -> list:
  7. """Returns the names of all immediate child folders of path.
  8. Args:
  9. path (str): path to search
  10. Returns:
  11. list: list of all child folder names
  12. """
  13. return [name for name in os.listdir(path) if os.path.isdir(os.path.join(path, name))]
  14. def list_jpegs_recursive(path: str) -> list:
  15. """Recursively lists all jpeg files in path.
  16. Args:
  17. path (str): path to search
  18. Returns:
  19. list: list of all jpeg files
  20. """
  21. print(os.path.join(path, "**/*.jpg"))
  22. return [name for name in glob(os.path.join(path, "**/*.jpg"), recursive=True) if os.path.isfile(os.path.join(path, name))]
  23. def verify_expected_subfolders(session_path: str):
  24. """Assert that the given session folder contains exactly the three subfolders Motion, Lapse, Full.
  25. Args:
  26. session_path (str): session folder path
  27. """
  28. subfolders = list_folders(session_path)
  29. if sorted(subfolders) != sorted(expected_subfolders):
  30. raise AssertionError(f"{session_path}: Expected subfolders {expected_subfolders} but found {subfolders}")
  31. # Pickle helpers
  32. def dump(filename: str, data):
  33. """Dumps data using pickle.
  34. Args:
  35. filename (str): Target file name.
  36. data (any): Data.
  37. """
  38. with open(filename, "wb") as f:
  39. pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
  40. def load(filename: str):
  41. """Loads data using pickle.
  42. Args:
  43. filename (str): Target file name.
  44. Returns:
  45. any: loaded data
  46. """
  47. with open(filename, "rb") as f:
  48. return pickle.load(f)