__init__.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. import os
  2. import warnings
  3. from modulefinder import Module
  4. import torch
  5. from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils
  6. from .extension import _HAS_OPS
  7. try:
  8. from .version import __version__ # noqa: F401
  9. except ImportError:
  10. pass
  11. # Check if torchvision is being imported within the root folder
  12. if not _HAS_OPS and os.path.dirname(os.path.realpath(__file__)) == os.path.join(
  13. os.path.realpath(os.getcwd()), "torchvision"
  14. ):
  15. message = (
  16. "You are importing torchvision within its own root folder ({}). "
  17. "This is not expected to work and may give errors. Please exit the "
  18. "torchvision project source and relaunch your python interpreter."
  19. )
  20. warnings.warn(message.format(os.getcwd()))
  21. _image_backend = "PIL"
  22. _video_backend = "pyav"
  23. def set_image_backend(backend):
  24. """
  25. Specifies the package used to load images.
  26. Args:
  27. backend (string): Name of the image backend. one of {'PIL', 'accimage'}.
  28. The :mod:`accimage` package uses the Intel IPP library. It is
  29. generally faster than PIL, but does not support as many operations.
  30. """
  31. global _image_backend
  32. if backend not in ["PIL", "accimage"]:
  33. raise ValueError(f"Invalid backend '{backend}'. Options are 'PIL' and 'accimage'")
  34. _image_backend = backend
  35. def get_image_backend():
  36. """
  37. Gets the name of the package used to load images
  38. """
  39. return _image_backend
  40. def set_video_backend(backend):
  41. """
  42. Specifies the package used to decode videos.
  43. Args:
  44. backend (string): Name of the video backend. one of {'pyav', 'video_reader'}.
  45. The :mod:`pyav` package uses the 3rd party PyAv library. It is a Pythonic
  46. binding for the FFmpeg libraries.
  47. The :mod:`video_reader` package includes a native C++ implementation on
  48. top of FFMPEG libraries, and a python API of TorchScript custom operator.
  49. It generally decodes faster than :mod:`pyav`, but is perhaps less robust.
  50. .. note::
  51. Building with FFMPEG is disabled by default in the latest `main`. If you want to use the 'video_reader'
  52. backend, please compile torchvision from source.
  53. """
  54. global _video_backend
  55. if backend not in ["pyav", "video_reader", "cuda"]:
  56. raise ValueError("Invalid video backend '%s'. Options are 'pyav', 'video_reader' and 'cuda'" % backend)
  57. if backend == "video_reader" and not io._HAS_VIDEO_OPT:
  58. # TODO: better messages
  59. message = "video_reader video backend is not available. Please compile torchvision from source and try again"
  60. raise RuntimeError(message)
  61. elif backend == "cuda" and not io._HAS_GPU_VIDEO_DECODER:
  62. # TODO: better messages
  63. message = "cuda video backend is not available."
  64. raise RuntimeError(message)
  65. else:
  66. _video_backend = backend
  67. def get_video_backend():
  68. """
  69. Returns the currently active video backend used to decode videos.
  70. Returns:
  71. str: Name of the video backend. one of {'pyav', 'video_reader'}.
  72. """
  73. return _video_backend
  74. def _is_tracing():
  75. return torch._C._get_tracing_state()
  76. def disable_beta_transforms_warning():
  77. # Noop, only exists to avoid breaking existing code.
  78. # See https://github.com/pytorch/vision/issues/7896
  79. pass