Migrate avif and heic decoders to torchvision-extra-decoders repo

Author: NicolasHug

torchvision/csrc/io/image/cpu/decode_image.cpp

@@ -50,6 +50,12 @@ torch::Tensor decode_image(
     return decode_gif(data);
   }
 
+  // TODO: We have to remove the "avif" detection logic from here, since
+  // `decode_avif` doesn't exist in torchvision anymore. We can reimplement the
+  // signature detection logic in Python instead. Same for heic. Longer term, we
+  // can move the logic to the torchvision-extra-decoders repo and rely on
+  // libmagic for a more accurate signature check.
+
   // We assume the signature of an avif file is
   // 0000 0020 6674 7970 6176 6966
   // xxxx xxxx  f t  y p  a v  i f

torchvision/io/image.py

@@ -377,6 +377,7 @@ def decode_webp(
     return torch.ops.image.decode_webp(input, mode.value)
 
 
+# TODO: remove this, and the associated avif decoder
 def _decode_avif(
     input: torch.Tensor,
     mode: ImageReadMode = ImageReadMode.UNCHANGED,
@@ -408,6 +409,7 @@ def _decode_avif(
     return torch.ops.image.decode_avif(input, mode.value)
 
 
+# TODO: remove this, and the associated heic decoder
 def _decode_heic(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     """
     Decode an HEIC image into a 3 dimensional RGB[A] Tensor.
@@ -434,3 +436,32 @@ def _decode_heic(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHAN
     if isinstance(mode, str):
         mode = ImageReadMode[mode.upper()]
     return torch.ops.image.decode_heic(input, mode.value)
+
+
+_EXTRA_DECODERS_ALREADY_LOADED = False
+
+
+def _load_extra_decoders_once():
+    global _EXTRA_DECODERS_ALREADY_LOADED
+    if _EXTRA_DECODERS_ALREADY_LOADED:
+        return
+
+    try:
+        import torchvision_extra_decoders
+    except ImportError as e:
+        raise RuntimeError("You need to pip install torchvision-extra-decoders blah blah blah") from e
+
+    # This will expose torch.ops.extra_decoders_ns.decode_avif and torch.ops.extra_decoders_ns.decode_heic
+    torchvision_extra_decoders.expose_extra_decoders()
+
+    _EXTRA_DECODERS_ALREADY_LOADED = True
+
+
+def decode_avif(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
+    _load_extra_decoders_once()
+    return torch.ops.extra_decoders_ns.decode_avif(input, mode.value)
+
+
+def decode_heic(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
+    _load_extra_decoders_once()
+    return torch.ops.extra_decoders_ns.decode_heic(input, mode.value)