Implement readScanline() for IvfAv1 decompressor

Summary: @public Similar to the `WebpDecompressor` where the underlying library does not inherently support reading scanline-by-scanline. The input is lazily loaded entirely when the first scanline is requested. Then those scanlines are stored and handed over as the outside world asks for more. Reviewed By: AurelC2G Differential Revision: D14893977 fbshipit-source-id: ae88ff7a901b1d0905d5058a8409120051d43afd
facebookincubator · Apr 16, 2019 · 16ed05a · 16ed05a
1 parent c46d7c8
commit 16ed05a
Show file tree

Hide file tree

Showing 3 changed files with 152 additions and 4 deletions.
diff --git a/cpp/spectrum/plugins/avif/IvfAv1Decompressor.cpp b/cpp/spectrum/plugins/avif/IvfAv1Decompressor.cpp
@@ -13,7 +13,9 @@
 
 #include <dav1d/dav1d.h>
 #include <folly/Optional.h>
+#include <folly/ScopeGuard.h>
 #include <ivf/ivfheader.h>
+#include <libyuv/convert_from.h>
 
 #include <memory>
 
@@ -53,7 +55,112 @@ IvfAv1Decompressor::~IvfAv1Decompressor() {
 // Private
 //
 
-void IvfAv1Decompressor::ensureHeaderIsRead() {
+void IvfAv1Decompressor::_ensureEntireImageIsRead() {
+  if (_entireImageHasBeenRead) {
+    return;
+  }
+  _entireImageHasBeenRead = true;
+
+  // read entire remaining image source
+  std::vector<std::uint8_t> payload(_payloadLength);
+  const auto bytesRead =
+      _source.read(reinterpret_cast<char*>(payload.data()), _payloadLength);
+  SPECTRUM_ERROR_CSTR_IF_NOT(
+      bytesRead == _payloadLength,
+      codecs::error::DecompressorFailure,
+      "actual AV1 payload was shorter than advertised in IVF header");
+
+  Dav1dData dav1dData{0};
+  SPECTRUM_ERROR_CSTR_IF_NOT(
+      DAV1D_OK ==
+          dav1d_data_wrap(
+              &dav1dData,
+              payload.data(),
+              payload.size(),
+              [](const uint8_t*, void*) { /* no-op free */ },
+              nullptr),
+      codecs::error::DecompressorFailure,
+      "failed dav1d_data_wrap");
+  SCOPE_EXIT {
+    dav1d_data_unref(&dav1dData);
+  };
+
+  SPECTRUM_ERROR_CSTR_IF_NOT(
+      DAV1D_OK == dav1d_send_data(_dav1dContext, &dav1dData),
+      codecs::error::DecompressorFailure,
+      "failed dav1d_send_data");
+
+  std::vector<std::uint8_t> rgbImage;
+  std::uint32_t rgbStride, totalRgbPixels, width, height;
+  {
+    Dav1dPicture dav1dPicture{0};
+    SPECTRUM_ERROR_CSTR_IF_NOT(
+        DAV1D_OK == dav1d_get_picture(_dav1dContext, &dav1dPicture),
+        codecs::error::DecompressorFailure,
+        "failed dav1d_get_picture");
+    SCOPE_EXIT {
+      dav1d_picture_unref(&dav1dPicture);
+    };
+
+    SPECTRUM_ERROR_CSTR_IF_NOT(
+        dav1dPicture.p.bpc == 8,
+        codecs::error::DecompressorFailure,
+        "can only read 8-bit images");
+
+    SPECTRUM_ERROR_CSTR_IF_NOT(
+        dav1dPicture.p.layout == DAV1D_PIXEL_LAYOUT_I420,
+        codecs::error::DecompressorFailure,
+        "can only read I420 images");
+
+    width = dav1dPicture.p.w;
+    height = dav1dPicture.p.h;
+    SPECTRUM_ERROR_CSTR_IF_NOT(
+        _imageSpecification->size.width == width &&
+            _imageSpecification->size.height == height,
+        codecs::error::DecompressorFailure,
+        "payload dimensions do not match IVF information");
+
+    rgbStride = dav1dPicture.p.w * 3;
+    totalRgbPixels = rgbStride * dav1dPicture.p.h;
+    rgbImage.resize(totalRgbPixels);
+
+    libyuv::I420ToRAW(
+        reinterpret_cast<uint8_t*>(dav1dPicture.data[0]), // Y
+        dav1dPicture.stride[0],
+        reinterpret_cast<uint8_t*>(dav1dPicture.data[1]), // U
+        dav1dPicture.stride[1],
+        reinterpret_cast<uint8_t*>(dav1dPicture.data[2]), // V
+        dav1dPicture.stride[1], // not an index mistake: 0=luma, 1=chroma
+        rgbImage.data(),
+        rgbStride,
+        dav1dPicture.p.w,
+        dav1dPicture.p.h);
+
+    // dav1dPicture will go out of scope here and we continue processing from
+    // argbImage only; this reduces the peak memory consumption
+  }
+
+  _entireImage.reserve(height);
+
+  for (auto row = 0; row < height; row++) {
+    auto scanline = std::make_unique<image::Scanline>(
+        image::pixel::specifications::RGB, width);
+
+    const auto rgbImageOffset = row * rgbStride;
+    std::copy_n(
+        rgbImage.data() + rgbImageOffset,
+        scanline->sizeBytes(),
+        scanline->data());
+
+    _entireImage.push_back(std::move(scanline));
+  }
+
+  // free context early
+  dav1d_close(&_dav1dContext);
+  _dav1dContext = nullptr;
+}
+
+void IvfAv1Decompressor::_ensureHeaderIsRead() {
   if (_imageSpecification.hasValue()) {
     return;
   }
@@ -69,6 +176,7 @@ void IvfAv1Decompressor::ensureHeaderIsRead() {
   // Read IvfFrameHeader
   const auto frameHeader = fb::ivf::parseIvfFrameHeader(buffer);
   SPECTRUM_ENFORCE_IF_NOT(frameHeader.payloadLength > 0);
+  _payloadLength = frameHeader.payloadLength;
 
   _imageSpecification = image::Specification{
       .size = image::Size{fileHeader.frameWidth, fileHeader.frameHeight},
@@ -81,7 +189,7 @@ void IvfAv1Decompressor::ensureHeaderIsRead() {
 //
 
 image::Specification IvfAv1Decompressor::sourceImageSpecification() {
-  ensureHeaderIsRead();
+  _ensureHeaderIsRead();
   return *_imageSpecification;
 }
 
@@ -90,7 +198,20 @@ image::Specification IvfAv1Decompressor::outputImageSpecification() {
 }
 
 std::unique_ptr<image::Scanline> IvfAv1Decompressor::readScanline() {
-  return nullptr;
+  _ensureEntireImageIsRead();
+  std::unique_ptr<image::Scanline> result;
+
+  if (_currentOutputScanline < _imageSpecification->size.height) {
+    result = std::move(_entireImage[_currentOutputScanline]);
+    _currentOutputScanline++;
+  }
+
+  // free memory early
+  if (_currentOutputScanline == _imageSpecification->size.height) {
+    _entireImage.clear();
+  }
+
+  return result;
 }
 
 } // namespace avif

diff --git a/cpp/spectrum/plugins/avif/IvfAv1Decompressor.h b/cpp/spectrum/plugins/avif/IvfAv1Decompressor.h
@@ -36,7 +36,13 @@ class IvfAv1Decompressor final : public codecs::IDecompressor {
   io::IImageSource& _source;
 
   folly::Optional<image::Specification> _imageSpecification;
-  void ensureHeaderIsRead();
+  size_t _payloadLength;
+  void _ensureHeaderIsRead();
+
+  bool _entireImageHasBeenRead = false;
+  std::vector<std::unique_ptr<image::Scanline>> _entireImage;
+  std::uint32_t _currentOutputScanline = 0;
+  void _ensureEntireImageIsRead();
 
  public:
   image::Specification sourceImageSpecification() override;

diff --git a/cpp/test/unit/plugins/avif/IvfAv1DecompressorTest.cpp b/cpp/test/unit/plugins/avif/IvfAv1DecompressorTest.cpp
@@ -39,6 +39,27 @@ TEST(
   ASSERT_EQ((image::Size{256, 170}), specifications.size);
 }
 
+TEST(
+    plugins_avif_IvfAv1Decompressor,
+    whenReadingImage_thenSpecsOfScanlinesReturnedMatchExpectations) {
+  io::FileImageSource source{
+      testdata::paths::avif::s256_170_rav1e_s420.normalized()};
+  auto decompressor = IvfAv1Decompressor{source};
+
+  const auto specifications = decompressor.sourceImageSpecification();
+  ASSERT_EQ((image::Size{256, 170}), specifications.size);
+
+  for (auto row = 0; row < specifications.size.height; row++) {
+    auto scanline = decompressor.readScanline();
+    ASSERT_TRUE(scanline != nullptr);
+    ASSERT_EQ(image::pixel::specifications::RGB, scanline->specification());
+    ASSERT_EQ(specifications.size.width, scanline->width());
+  }
+
+  auto scanline = decompressor.readScanline();
+  ASSERT_TRUE(scanline == nullptr);
+}
+
 } // namespace test
 } // namespace avif
 } // namespace plugins