|
| 1 | +#include "decode_avif.h" |
| 2 | + |
| 3 | +#if AVIF_FOUND |
| 4 | +#include "avif/avif.h" |
| 5 | +#endif // AVIF_FOUND |
| 6 | + |
| 7 | +namespace vision { |
| 8 | +namespace image { |
| 9 | + |
| 10 | +#if !AVIF_FOUND |
| 11 | +torch::Tensor decode_avif(const torch::Tensor& data) { |
| 12 | + TORCH_CHECK( |
| 13 | + false, "decode_avif: torchvision not compiled with libavif support"); |
| 14 | +} |
| 15 | +#else |
| 16 | + |
| 17 | +// This normally comes from avif_cxx.h, but it's not always present when |
| 18 | +// installing libavif. So we just copy/paste it here. |
| 19 | +struct UniquePtrDeleter { |
| 20 | + void operator()(avifDecoder* decoder) const { |
| 21 | + avifDecoderDestroy(decoder); |
| 22 | + } |
| 23 | +}; |
| 24 | +using DecoderPtr = std::unique_ptr<avifDecoder, UniquePtrDeleter>; |
| 25 | + |
| 26 | +torch::Tensor decode_avif(const torch::Tensor& encoded_data) { |
| 27 | + // This is based on |
| 28 | + // https://github.com/AOMediaCodec/libavif/blob/main/examples/avif_example_decode_memory.c |
| 29 | + // Refer there for more detail about what each function does, and which |
| 30 | + // structure/data is available after which call. |
| 31 | + |
| 32 | + TORCH_CHECK(encoded_data.is_contiguous(), "Input tensor must be contiguous."); |
| 33 | + TORCH_CHECK( |
| 34 | + encoded_data.dtype() == torch::kU8, |
| 35 | + "Input tensor must have uint8 data type, got ", |
| 36 | + encoded_data.dtype()); |
| 37 | + TORCH_CHECK( |
| 38 | + encoded_data.dim() == 1, |
| 39 | + "Input tensor must be 1-dimensional, got ", |
| 40 | + encoded_data.dim(), |
| 41 | + " dims."); |
| 42 | + |
| 43 | + DecoderPtr decoder(avifDecoderCreate()); |
| 44 | + TORCH_CHECK(decoder != nullptr, "Failed to create avif decoder."); |
| 45 | + |
| 46 | + auto result = AVIF_RESULT_UNKNOWN_ERROR; |
| 47 | + result = avifDecoderSetIOMemory( |
| 48 | + decoder.get(), encoded_data.data_ptr<uint8_t>(), encoded_data.numel()); |
| 49 | + TORCH_CHECK( |
| 50 | + result == AVIF_RESULT_OK, |
| 51 | + "avifDecoderSetIOMemory failed:", |
| 52 | + avifResultToString(result)); |
| 53 | + |
| 54 | + result = avifDecoderParse(decoder.get()); |
| 55 | + TORCH_CHECK( |
| 56 | + result == AVIF_RESULT_OK, |
| 57 | + "avifDecoderParse failed: ", |
| 58 | + avifResultToString(result)); |
| 59 | + TORCH_CHECK( |
| 60 | + decoder->imageCount == 1, "Avif file contains more than one image"); |
| 61 | + TORCH_CHECK( |
| 62 | + decoder->image->depth <= 8, |
| 63 | + "avif images with bitdepth > 8 are not supported"); |
| 64 | + |
| 65 | + result = avifDecoderNextImage(decoder.get()); |
| 66 | + TORCH_CHECK( |
| 67 | + result == AVIF_RESULT_OK, |
| 68 | + "avifDecoderNextImage failed:", |
| 69 | + avifResultToString(result)); |
| 70 | + |
| 71 | + auto out = torch::empty( |
| 72 | + {decoder->image->height, decoder->image->width, 3}, torch::kUInt8); |
| 73 | + |
| 74 | + avifRGBImage rgb; |
| 75 | + memset(&rgb, 0, sizeof(rgb)); |
| 76 | + avifRGBImageSetDefaults(&rgb, decoder->image); |
| 77 | + rgb.format = AVIF_RGB_FORMAT_RGB; |
| 78 | + rgb.pixels = out.data_ptr<uint8_t>(); |
| 79 | + rgb.rowBytes = rgb.width * avifRGBImagePixelSize(&rgb); |
| 80 | + |
| 81 | + result = avifImageYUVToRGB(decoder->image, &rgb); |
| 82 | + TORCH_CHECK( |
| 83 | + result == AVIF_RESULT_OK, |
| 84 | + "avifImageYUVToRGB failed: ", |
| 85 | + avifResultToString(result)); |
| 86 | + |
| 87 | + return out.permute({2, 0, 1}); // return CHW, channels-last |
| 88 | +} |
| 89 | +#endif // AVIF_FOUND |
| 90 | + |
| 91 | +} // namespace image |
| 92 | +} // namespace vision |
0 commit comments