Hardware decoding of frames (#13)

Allow user to initialize FrameDecoder with a HW device type selected. Performs decoding (but not pixel format conversion and scaling) on the hardware device.

Signed-off-by: Emerson Knapp <emerson.b.knapp@gmail.com>

Author: emersonknapp

broll/CMakeLists.txt

@@ -78,6 +78,9 @@ if(BUILD_TESTING)
   ament_lint_auto_find_test_dependencies()
 endif()
 
+ament_export_dependencies(
+  sensor_msgs
+)
 ament_export_include_directories(
   "include/${PROJECT_NAME}"
 )

broll/include/broll/frame_decoder.hpp

@@ -38,11 +38,15 @@ class FrameDecoder
   /// @param target_fmt Pixel format to convert to, if necessary.
   ///   AV_PIX_FMT_NONE guarantees no conversion.
   /// @param scale Scale applied to image dimensions (by multiplication, 0.5 is half size)
+  /// @param hw_device_type Attempt to use hardware decoding device of this type.
+  ///   AV_HWDEVICE_TYPE_NONE uses software decoding only.
   /// @param dbg_print Print info about each decoded frame to stdout
+  /// @throws std::invalid_argument if hw_device_type is not supported
   FrameDecoder(
     AVCodecID codec_id,
     AVPixelFormat target_fmt = AV_PIX_FMT_NONE,
     double scale = 1.0f,
+    AVHWDeviceType hw_device_type = AV_HWDEVICE_TYPE_NONE,
     bool dbg_print = false);
   virtual ~FrameDecoder();
 
@@ -79,6 +83,11 @@ class FrameDecoder
   /// available in HEVC decoder internals, so we would need to use libx265 directly to detect.
   void startSkippingPFrames();
 
+  /// @brief Function to help AVCodecContext pick a pixel format that matches hwPixFmt_.
+  ///
+  /// Only used if hardware decoding is enabled.
+  static AVPixelFormat getHardwarePixelFormat(AVCodecContext * ctx, const AVPixelFormat * pix_fmts);
+
   AVPacket * packet_ = nullptr;
   const AVCodec * codec_ = nullptr;
   AVCodecContext * codecCtx_ = nullptr;
@@ -87,6 +96,17 @@ class FrameDecoder
   AVFrame * decodedFrame_ = nullptr;
   AVFrame * convertedFrame_ = nullptr;
 
+  // Hardware decoding extras
+  // Pixel format reported for an on-hardware DMA frame, such as "cuda", which isn't
+  // a real pixel format but how the frame reports that it isn't in CPU memory.
+  AVPixelFormat hwPixFmt_ = AV_PIX_FMT_NONE;
+  // Pixel format to convert to when transferring the frame from hardware into CPU memory.
+  // Hardware devices offer several options, but not all that we might want so this
+  // is the pre-conversion format before sws.
+  AVPixelFormat hwSoftwarePixFmt_ = AV_PIX_FMT_NONE;
+  AVBufferRef * hwDeviceCtx_ = nullptr;
+  AVFrame * hwFrame_ = nullptr;
+
   float scale_ = 1.0f;
   uint scaled_width_ = 0u;
   uint scaled_height_ = 0u;

broll/src/frame_decoder.cpp

@@ -44,13 +44,11 @@ static AVFrame * allocPicture(enum AVPixelFormat pix_fmt, int width, int height)
   int ret = av_frame_get_buffer(picture, 0);
   if (ret < 0) {
     fprintf(stderr, "Could not allocate frame data.\n");
-    exit(1);
+    return nullptr;
   }
   return picture;
 }
 
-static const int64_t NS_TO_S = 1000000000;
-
 }  // namespace
 
 namespace broll
@@ -60,6 +58,7 @@ FrameDecoder::FrameDecoder(
   AVCodecID codec_id,
   AVPixelFormat target_fmt,
   double scale,
+  AVHWDeviceType hw_device_type,
   bool dbg_print)
 : targetPixFmt_(target_fmt),
   scale_(scale),
@@ -91,6 +90,61 @@ FrameDecoder::FrameDecoder(
     assert(false && "failed to copy codec params to codec context");
   }
 
+  if (hw_device_type != AV_HWDEVICE_TYPE_NONE) {
+    // Read HW configs 0->N until finding the one we're looking for or it returns nullptr
+    for (size_t i = 0;; i++) {
+      const AVCodecHWConfig * config = avcodec_get_hw_config(codec_, i);
+      if (!config) {
+        BROLL_LOG_ERROR(
+          "Decoder %s does not support device type %s.\n",
+          codec_->name, av_hwdevice_get_type_name(hw_device_type));
+        throw std::runtime_error("Unsupported hardware device type");
+      }
+      if (
+        config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX &&
+        config->device_type == hw_device_type)
+      {
+        hwPixFmt_ = config->pix_fmt;
+        break;
+      }
+    }
+
+    BROLL_LOG_INFO(
+      "Hardware decoding enabled. Pixel format \"%s\"",
+      av_get_pix_fmt_name(hwPixFmt_));
+    codecCtx_->get_format = FrameDecoder::getHardwarePixelFormat;
+    int err = av_hwdevice_ctx_create(&hwDeviceCtx_, hw_device_type, nullptr, nullptr, 0);
+    if (err < 0) {
+      assert(false && "Failed to initialize hardware decoder.");
+    }
+    codecCtx_->hw_device_ctx = av_buffer_ref(hwDeviceCtx_);
+    {
+      // Check what software-side pixel formats the hardware transfer supports.
+      // Print all for visibility, but select just the first one.
+      // This should probably be configurable, to let user target best pre-sws fmt.
+      AVHWFramesConstraints * hw_frames_const = av_hwdevice_get_hwframe_constraints(
+        hwDeviceCtx_, nullptr);
+      assert(hw_frames_const && "Couldn't retrieve hardware device frame constraints.");
+      hwSoftwarePixFmt_ = AV_PIX_FMT_NONE;
+      for (AVPixelFormat * p = hw_frames_const->valid_sw_formats; *p != AV_PIX_FMT_NONE; p++) {
+        if (sws_isSupportedInput(*p)) {
+          BROLL_LOG_INFO(
+            "Supported hardware-to-software pixel format: \"%s\"",
+            av_get_pix_fmt_name(*p));
+          if (hwSoftwarePixFmt_ == AV_PIX_FMT_NONE) {
+            hwSoftwarePixFmt_ = *p;
+          }
+        }
+      }
+      av_hwframe_constraints_free(&hw_frames_const);
+      BROLL_LOG_INFO("Selected hw/sw pixel format: \"%s\"", av_get_pix_fmt_name(hwSoftwarePixFmt_));
+    }
+
+    BROLL_LOG_INFO("Succeeded to init hardware decoder");
+    hwFrame_ = av_frame_alloc();
+    assert(hwFrame_ && "failed to alloc hardwareFrame");
+  }
+
   if (avcodec_open2(codecCtx_, codec_, nullptr) < 0) {
     assert(false && "failed to open codec through avcodec_open2");
   }
@@ -122,7 +176,13 @@ bool FrameDecoder::decodeFrame(const AVPacket & packet_in, AVFrame & frame_out)
     BROLL_LOG_ERROR("avcodec_send_packet failed: %s", errStr);
     return false;
   }
-  const int recv_frame_resp = avcodec_receive_frame(codecCtx_, &frame_out);
+
+  int recv_frame_resp;
+  if (hwFrame_) {
+    recv_frame_resp = avcodec_receive_frame(codecCtx_, hwFrame_);
+  } else {
+    recv_frame_resp = avcodec_receive_frame(codecCtx_, &frame_out);
+  }
   if (recv_frame_resp == AVERROR(EAGAIN)) {
     BROLL_LOG_DEBUG("avcodec_receive_frame returned EAGAIN");
     return false;
@@ -137,6 +197,20 @@ bool FrameDecoder::decodeFrame(const AVPacket & packet_in, AVFrame & frame_out)
     return false;
   }
 
+  if (hwFrame_) {
+    if (hwFrame_->format != hwPixFmt_) {
+      BROLL_LOG_ERROR("Received hardware frame was not in expected pixel format.");
+      return false;
+    }
+    // Move the frame from hardware device memory to CPU memory.
+    // Setting the pixel format on frame_out prompts the transfer to convert to a specific format.
+    frame_out.format = hwSoftwarePixFmt_;
+    if (av_hwframe_transfer_data(&frame_out, hwFrame_, 0) < 0) {
+      BROLL_LOG_ERROR("Error transferring the data to system memory");
+      return false;
+    }
+  }
+
 #if LIBAVCODEC_VERSION_MAJOR >= 60
   int64_t frame_num = codecCtx_->frame_num;
   bool is_key_frame = frame_out.flags & AV_FRAME_FLAG_KEY;
@@ -231,13 +305,18 @@ bool FrameDecoder::decode(const AVPacket & in, sensor_msgs::msg::Image & out)
     BROLL_LOG_INFO(
       "Frame Decoder initialized: resolution in %d x %d, resolution out %d x %d",
       width, height, scaled_width_, scaled_height_);
-    BROLL_LOG_INFO("\tCodec %s ID %d", codec_->name, codec_->id);
+    BROLL_LOG_INFO("\tCodec %d ('%s')", codec_->id, codec_->name);
+    BROLL_LOG_INFO(
+      "\tCodec pixfmt '%s', decoded pixfmt '%s'",
+      av_get_pix_fmt_name(codecCtx_->pix_fmt),
+      av_get_pix_fmt_name(static_cast<AVPixelFormat>(decodedFrame_->format)));
+    BROLL_LOG_INFO("\tTarget pixfmt '%s'", av_get_pix_fmt_name(targetPixFmt_));
 
     convertedFrame_ = allocPicture(targetPixFmt_, scaled_width_, scaled_height_);
     assert(convertedFrame_ && "failed to alloc convertedFrame");
 
     bool set_extended_color_range = false;
-    AVPixelFormat sws_pix_fmt = codecCtx_->pix_fmt;
+    AVPixelFormat sws_pix_fmt = static_cast<AVPixelFormat>(decodedFrame_->format);
     switch (sws_pix_fmt) {
       case AV_PIX_FMT_YUVJ420P:
         sws_pix_fmt = AV_PIX_FMT_YUV420P;
@@ -299,4 +378,19 @@ bool FrameDecoder::decode(
   return res;
 }
 
+AVPixelFormat FrameDecoder::getHardwarePixelFormat(
+  AVCodecContext * ctx, const AVPixelFormat * pix_fmts)
+{
+  const AVPixelFormat * p;
+  const AVPixelFormat hwPixFmt = static_cast<FrameDecoder *>(ctx->opaque)->hwPixFmt_;
+  for (p = pix_fmts; *p != AV_PIX_FMT_NONE; p++) {
+    if (*p == hwPixFmt) {
+      return *p;
+    }
+  }
+  BROLL_LOG_ERROR("Failed to get HW surface format.");
+  return AV_PIX_FMT_NONE;
+}
+
+
 }  // namespace broll

rosbag2_storage_broll/CMakeLists.txt

@@ -78,6 +78,9 @@ if(BUILD_TESTING)
   ament_lint_auto_find_test_dependencies()
 endif()
 
+ament_export_dependencies(
+  rosbag2_transport
+)
 ament_export_libraries(
   bag_utils
 )