From 6c9c3ce838c70d50bd0cc392811cc68bbf8931ef Mon Sep 17 00:00:00 2001
From: Nicolas Hug <contact@nicolas-hug.com>
Date: Thu, 12 Sep 2024 15:17:33 +0100
Subject: [PATCH] Expose libavif and heic internally

---
 test/test_image.py                        | 10 ++++++++++
 torchvision/_internally_replaced_utils.py |  1 +
 torchvision/io/__init__.py                |  7 +++++++
 3 files changed, 18 insertions(+)

diff --git a/test/test_image.py b/test/test_image.py
index 624dc57fada..f3c2984b348 100644
--- a/test/test_image.py
+++ b/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__])
diff --git a/torchvision/_internally_replaced_utils.py b/torchvision/_internally_replaced_utils.py
index d9a6e261ea2..e0fa72489f1 100644
--- a/torchvision/_internally_replaced_utils.py
+++ b/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:
diff --git a/torchvision/io/__init__.py b/torchvision/io/__init__.py
index a604ea1fdb6..0dcbd7e9cea 100644
--- a/torchvision/io/__init__.py
+++ b/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"]