PixelFrame async processing helpers

Summary: When trying to speed up image decoding, we need to do the image processing in a separate than the disk/vrs thread. These new methods allow us to store the image data as read quickly from disk, so that image can be decoded and/or converted in another thread, reducing contention on the thread reading the vrs file.

Reviewed By: kiminoue7

Differential Revision: D51789222

fbshipit-source-id: 44019a7591749b0bebf4eda83b9faf27be02c51e

vrs/utils/PixelFrame.cpp

@@ -25,6 +25,7 @@
 #include <logging/Verify.h>
 
 #include <vrs/helpers/FileMacros.h>
+#include <vrs/utils/BufferRecordReader.hpp>
 #include <vrs/utils/converters/Raw10ToGrey10Converter.h>
 
 using namespace std;
@@ -179,16 +180,29 @@ PixelFrame::PixelFrame(const ImageContentBlockSpec& spec)
   }
 }
 
+PixelFrame::PixelFrame(const ImageContentBlockSpec& spec, vector<uint8_t>&& frameBytes)
+    : imageSpec_{spec}, frameBytes_{std::move(frameBytes)} {}
+
 void PixelFrame::init(const ImageContentBlockSpec& spec) {
-  if (!hasSamePixels(spec)) {
-    imageSpec_ = spec;
+  if (imageSpec_.getImageFormat() != ImageFormat::RAW || !hasSamePixels(spec)) {
+    imageSpec_ = {
+        spec.getPixelFormat(),
+        spec.getWidth(),
+        spec.getHeight(),
+        spec.getRawStride(),
+        spec.getRawStride2()};
     size_t size = imageSpec_.getRawImageSize();
     if (XR_VERIFY(size != ContentBlock::kSizeUnknown)) {
       frameBytes_.resize(size);
     }
   }
 }
 
+void PixelFrame::init(const ImageContentBlockSpec& spec, vector<uint8_t>&& frameBytes) {
+  imageSpec_ = spec;
+  frameBytes_ = std::move(frameBytes);
+}
+
 void PixelFrame::init(shared_ptr<PixelFrame>& inOutFrame, const ImageContentBlockSpec& spec) {
   if (!inOutFrame) {
     inOutFrame = make_shared<PixelFrame>(spec);
@@ -260,6 +274,60 @@ bool PixelFrame::readFrame(RecordReader* reader, const ContentBlock& cb) {
   return false;
 }
 
+bool PixelFrame::readDiskImageData(
+    shared_ptr<PixelFrame>& frame,
+    RecordReader* reader,
+    const ContentBlock& cb) {
+  if (!frame) {
+    frame = make_shared<PixelFrame>();
+  }
+  return frame->readDiskImageData(reader, cb);
+}
+
+bool PixelFrame::readDiskImageData(RecordReader* reader, const ContentBlock& cb) {
+  size_t blockSize = cb.getBlockSize();
+  if (cb.getContentType() != ContentType::IMAGE || blockSize == ContentBlock::kSizeUnknown) {
+    return false;
+  }
+  const auto& spec = cb.image();
+  if (spec.getImageFormat() == ImageFormat::RAW) {
+    return readRawFrame(reader, spec);
+  }
+  imageSpec_ = spec;
+  frameBytes_.resize(blockSize);
+  return VERIFY_SUCCESS(reader->read(frameBytes_.data(), blockSize));
+}
+
+bool PixelFrame::decompressImage(VideoFrameHandler* videoFrameHandler) {
+  switch (imageSpec_.getImageFormat()) {
+    case ImageFormat::RAW:
+      return true;
+    case ImageFormat::VIDEO:
+      if (videoFrameHandler != nullptr) {
+        vector<uint8_t> compressedData(std::move(frameBytes_));
+        BufferReader reader;
+        return videoFrameHandler->tryToDecodeFrame(
+                   *this, reader.init(compressedData), {imageSpec_, compressedData.size()}) == 0;
+      }
+      break;
+    case ImageFormat::PNG: {
+      vector<uint8_t> compressedData(std::move(frameBytes_));
+      return readPngFrame(compressedData);
+    }
+    case ImageFormat::JPG: {
+      vector<uint8_t> compressedData(std::move(frameBytes_));
+      return readJpegFrame(compressedData);
+    }
+    case ImageFormat::JXL: {
+      vector<uint8_t> compressedData(std::move(frameBytes_));
+      return readJxlFrame(compressedData);
+    }
+    default:
+      return false;
+  }
+  return false;
+}
+
 bool PixelFrame::readRawFrame(RecordReader* reader, const ImageContentBlockSpec& inputImageSpec) {
   // Read multiplane images as is
   if (inputImageSpec.getPlaneCount() != 1) {

vrs/utils/PixelFrame.h

@@ -50,13 +50,15 @@ class PixelFrame {
  public:
   PixelFrame() = default;
   PixelFrame(const ImageContentBlockSpec& spec);
+  PixelFrame(const ImageContentBlockSpec& spec, vector<uint8_t>&& frameBytes);
   PixelFrame(PixelFormat pf, uint32_t w, uint32_t h, uint32_t stride = 0)
       : PixelFrame(ImageContentBlockSpec(pf, w, h, stride)) {}
 
   void init(const ImageContentBlockSpec& spec);
   inline void init(PixelFormat pf, uint32_t w, uint32_t h, uint32_t stride = 0) {
     init(ImageContentBlockSpec(pf, w, h, stride));
   }
+  void init(const ImageContentBlockSpec& spec, vector<uint8_t>&& frameBytes);
 
   static void init(shared_ptr<PixelFrame>& inOutFrame, const ImageContentBlockSpec& spec);
   static inline void init(
@@ -118,11 +120,22 @@ class PixelFrame {
   readFrame(shared_ptr<PixelFrame>& frame, RecordReader* reader, const ContentBlock& cb);
   bool readFrame(RecordReader* reader, const ContentBlock& cb);
 
+  /// Read a record's image data, merely reading the disk data without any decompression.
+  /// The resulting PixelFrame will have an unmodified ImageFormat (raw, jpg, png, jxl, video).
+  static bool
+  readDiskImageData(shared_ptr<PixelFrame>& frame, RecordReader* reader, const ContentBlock& cb);
+  bool readDiskImageData(RecordReader* reader, const ContentBlock& cb);
+
+  /// From any ImageFormat, decompress the image to ImageFormat::RAW if necessary.
+  /// To decompress ImageFormat::VIDEO data, you must provide a valid VideoFrameHandler object, the
+  /// same one for all the frames of a particular stream.
+  bool decompressImage(VideoFrameHandler* videoFrameHandler = nullptr);
+
   /// Read a RAW frame into the internal buffer.
   /// @return True if the frame type is supported & the frame was read.
   bool readRawFrame(RecordReader* reader, const ImageContentBlockSpec& inputImageSpec);
 
-  /// Read a compressed image, except for video codec compression.
+  /// Decode compressed image data, except for video codec compression.
   bool readCompressedFrame(const vector<uint8_t>& pixels, ImageFormat imageFormat);
 
   static bool readRawFrame(