__init__.py 3.0 KB

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