__init__.py 2.09 KB
Newer Older
1
2
import warnings

3
4
from torchvision import models
from torchvision import datasets
5
from torchvision import ops
6
7
from torchvision import transforms
from torchvision import utils
8
from torchvision import io
9

10
from .extension import _HAS_OPS
eellison's avatar
eellison committed
11
import torch
12

13
14
try:
    from .version import __version__  # noqa: F401
Soumith Chintala's avatar
Soumith Chintala committed
15
except ImportError:
16
    pass
17
18
19

_image_backend = 'PIL'

20
21
_video_backend = "pyav"

22
23
24
25
26
27

def set_image_backend(backend):
    """
    Specifies the package used to load images.

    Args:
28
29
30
        backend (string): Name of the image backend. one of {'PIL', 'accimage'}.
            The :mod:`accimage` package uses the Intel IPP library. It is
            generally faster than PIL, but does not support as many operations.
31
32
33
34
35
36
37
38
39
40
41
42
43
    """
    global _image_backend
    if backend not in ['PIL', 'accimage']:
        raise ValueError("Invalid backend '{}'. Options are 'PIL' and 'accimage'"
                         .format(backend))
    _image_backend = backend


def get_image_backend():
    """
    Gets the name of the package used to load images
    """
    return _image_backend
44
45


46
47
48
49
50
51
52
def set_video_backend(backend):
    """
    Specifies the package used to decode videos.

    Args:
        backend (string): Name of the video backend. one of {'pyav', 'video_reader'}.
            The :mod:`pyav` package uses the 3rd party PyAv library. It is a Pythonic
53
54
55
56
            binding for the FFmpeg libraries.
            The :mod:`video_reader` package includes a native C++ implementation on
            top of FFMPEG libraries, and a python API of TorchScript custom operator.
            It is generally decoding faster than :mod:`pyav`, but perhaps is less robust.
57
58
59
60
61
62
    """
    global _video_backend
    if backend not in ["pyav", "video_reader"]:
        raise ValueError(
            "Invalid video backend '%s'. Options are 'pyav' and 'video_reader'" % backend
        )
63
64
65
66
    if backend == "video_reader" and not io._HAS_VIDEO_OPT:
        warnings.warn("video_reader video backend is not available")
    else:
        _video_backend = backend
67
68
69
70
71
72


def get_video_backend():
    return _video_backend


73
74
def _is_tracing():
    return torch._C._get_tracing_state()