Expose libavif and heic internally (#8647)

test/test_image.py

@@ -4,6 +4,7 @@
 import os
 import re
 import sys
+from contextlib import nullcontext
 from pathlib import Path
 
 import numpy as np
@@ -13,6 +14,7 @@
 import torchvision.transforms.v2.functional as F
 from common_utils import assert_equal, cpu_and_cuda, IN_OSS_CI, needs_cuda
 from PIL import __version__ as PILLOW_VERSION, Image, ImageOps, ImageSequence
+from torchvision._internally_replaced_utils import IN_FBCODE
 from torchvision.io.image import (
  _decode_avif,
  _decode_heic,
@@ -1076,5 +1078,13 @@ def test_mode_str():
  assert decode_image(path, mode="RGBA").shape[0] == 4
 
 
+def test_avif_heic_fbcode():
+ cm = nullcontext() if IN_FBCODE else pytest.raises(ImportError, match="cannot import")
+ with cm:
+ from torchvision.io import decode_heic # noqa
+ with cm:
+ from torchvision.io import decode_avif # noqa
+
+
 if __name__ == "__main__":
  pytest.main([__file__])

torchvision/_internally_replaced_utils.py

@@ -6,6 +6,7 @@
 
 _HOME = os.path.join(_get_torch_home(), "datasets", "vision")
 _USE_SHARDED_DATASETS = False
+IN_FBCODE = False
 
 
 def _download_file_from_remote_location(fpath: str, url: str) -> None:

torchvision/io/__init__.py

@@ -74,3 +74,10 @@
  "Video",
  "VideoReader",
 ]
+
+from .._internally_replaced_utils import IN_FBCODE
+
+if IN_FBCODE:
+ from .image import _decode_avif as decode_avif, _decode_heic as decode_heic
+
+ __all__ += ["decode_avif", "decode_heic"]