__init__.py 2.1 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
11
from .extension import _HAS_OPS

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

_image_backend = 'PIL'

19
20
_video_backend = "pyav"

21
22
23
24
25
26

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

    Args:
27
28
29
        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.
30
31
32
33
34
35
36
37
38
39
40
41
42
    """
    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
43
44


45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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
                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 pyav, but perhaps is less robust.
    """
    global _video_backend
    if backend not in ["pyav", "video_reader"]:
        raise ValueError(
            "Invalid video backend '%s'. Options are 'pyav' and 'video_reader'" % backend
        )
62
63
64
65
    if backend == "video_reader" and not io._HAS_VIDEO_OPT:
        warnings.warn("video_reader video backend is not available")
    else:
        _video_backend = backend
66
67
68
69
70
71


def get_video_backend():
    return _video_backend


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