Program Listing for File encoder.hpp

↰ Return to documentation for file (include/ffmpeg_encoder_decoder/encoder.hpp)

// -*-c++-*---------------------------------------------------------------------------------------
// Copyright 2024 Bernd Pfrommer <bernd.pfrommer@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef FFMPEG_ENCODER_DECODER__ENCODER_HPP_
#define FFMPEG_ENCODER_DECODER__ENCODER_HPP_

#include <ffmpeg_encoder_decoder/pts_map.hpp>
#include <ffmpeg_encoder_decoder/tdiff.hpp>
#include <ffmpeg_encoder_decoder/types.hpp>
#include <functional>
#include <iostream>
#include <list>
#include <memory>
#include <mutex>
#include <opencv2/core/core.hpp>
#include <rclcpp/rclcpp.hpp>
#include <unordered_map>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavformat/avio.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
#include <libavutil/samplefmt.h>
#include <libswscale/swscale.h>
}

namespace ffmpeg_encoder_decoder
{
class Encoder
{
public:
  using Callback = std::function<void(
    const std::string & frame_id, const rclcpp::Time & stamp, const std::string & codec,
    uint32_t width, uint32_t height, uint64_t pts, uint8_t flags, uint8_t * data, size_t sz)>;

  Encoder();
  ~Encoder();
  void setEncoder(const std::string & encoder_name);
  void addAVOption(const std::string & key, const std::string & value)
  {
    Lock lock(mutex_);
    avOptions_.push_back({key, value});
  }

  void setAVSourcePixelFormat(const std::string & fmt)
  {
    Lock lock(mutex_);
    pixFormat_ = pixelFormat(fmt);
  }
  void setCVBridgeTargetFormat(const std::string & fmt)
  {
    Lock lock(mutex_);
    cvBridgeTargetFormat_ = fmt;
  }

  void setQMax(int q)
  {
    Lock lock(mutex_);
    qmax_ = q;
  }
  void setBitRate(int bitrate)
  {
    Lock lock(mutex_);
    bitRate_ = bitrate;
  }
  int getBitRate() const
  {
    Lock lock(mutex_);
    return (bitRate_);
  }

  void setGOPSize(int g)
  {
    Lock lock(mutex_);
    GOPSize_ = g;
  }
  int getGOPSize() const
  {
    Lock lock(mutex_);
    return (GOPSize_);
  }
  void setMaxBFrames(int b)
  {
    Lock lock(mutex_);
    maxBFrames_ = b;
  }
  int getMaxBFrames() const
  {
    Lock lock(mutex_);
    return (maxBFrames_);
  }
  void setFrameRate(int frames, int second)
  {
    Lock lock(mutex_);
    frameRate_.num = frames;
    frameRate_.den = second;
    timeBase_.num = second;
    timeBase_.den = frames;
  }
  void setMeasurePerformance(bool p)
  {
    Lock lock(mutex_);
    measurePerformance_ = p;
  }
  // ------- related to teardown and startup
  bool isInitialized() const
  {
    Lock lock(mutex_);
    return (codecContext_ != NULL);
  }
  bool initialize(
    int width, int height, Callback callback, const std::string & encoding = std::string());
  void setLogger(rclcpp::Logger logger) { logger_ = logger; }
  void reset();
  void encodeImage(const Image & msg);
  void flush();
  static std::string findCodec(const std::string & encoder);

  // ------- performance statistics
  void printTimers(const std::string & prefix) const;
  void resetTimers();
  // ------- deprecated functions
  [[deprecated("use addAVOption() instead.")]] void setProfile(const std::string & p)
  {
    addAVOption("profile", p);
  }
  [[deprecated("use addAVOption() instead.")]] void setPreset(const std::string & p)
  {
    addAVOption("preset", p);
  }
  [[deprecated("use addAVOption() instead.")]] void setTune(const std::string & p)
  {
    addAVOption("tune", p);
  }
  [[deprecated("use addAVOption() instead.")]] void setDelay(const std::string & p)
  {
    addAVOption("delay", p);
  }
  [[deprecated("use addAVOption() instead.")]] void setCRF(const std::string & c)
  {
    addAVOption("crf", c);
  }
  [[deprecated("use setAVSourcePixelFormat(fmt) instead.")]] void setPixelFormat(
    const std::string & fmt)
  {
    setAVSourcePixelFormat(fmt);
  }
  [[deprecated("use encodeImage(const Image &msg) instead.")]] void encodeImage(
    const cv::Mat & img, const Header & header, const rclcpp::Time & t0);
  [[deprecated("use flush() instead.")]] void flush(const Header & header);

private:
  using Lock = std::unique_lock<std::recursive_mutex>;

  bool openCodec(int width, int height, const std::string & encoding);
  void doOpenCodec(int width, int height, const std::string & encoding);
  void closeCodec();
  void doEncodeImage(const cv::Mat & img, const Header & header, const rclcpp::Time & t0);
  int drainPacket(int width, int height);
  AVPixelFormat pixelFormat(const std::string & f) const;
  void openHardwareDevice(
    const AVCodec * codec, enum AVHWDeviceType hwDevType, int width, int height);
  void setAVOption(const std::string & field, const std::string & value);
  enum AVPixelFormat findMatchingSourceFormat(
    const std::string & rosSrcFormat, enum AVPixelFormat targetFormat);

  // --------- variables
  rclcpp::Logger logger_;
  mutable std::recursive_mutex mutex_;
  Callback callback_;
  // config
  std::string encoder_;   // e.g. "libx264"
  std::string codec_;     // e.g. "h264"
  std::string encoding_;  // e.g. "h264/rgb8"
  int qmax_{-1};          // max allowed quantization. The lower the better quality
  int GOPSize_{-1};       // distance between two keyframes
  int maxBFrames_{-1};    // maximum number of b-frames
  int64_t bitRate_{0};    // max rate in bits/s
  std::vector<std::pair<std::string, std::string>> avOptions_;

  AVPixelFormat pixFormat_{AV_PIX_FMT_NONE};
  std::string cvBridgeTargetFormat_ = "bgr8";
  AVRational timeBase_{1, 100};
  AVRational frameRate_{100, 1};
  bool usesHardwareFrames_{false};
  std::string avSourcePixelFormat_;
  // ------ libav state
  AVCodecContext * codecContext_{nullptr};
  AVBufferRef * hwDeviceContext_{nullptr};
  AVFrame * frame_{nullptr};
  AVFrame * hw_frame_{nullptr};
  AVPacket * packet_{nullptr};
  // ------ libswscale state
  AVFrame * wrapperFrame_{nullptr};
  SwsContext * swsContext_{NULL};
  // ---------- other stuff
  int64_t pts_{0};
  PTSMap ptsToStamp_;
  // ---------- performance analysis
  bool measurePerformance_{true};
  int64_t totalInBytes_{0};
  int64_t totalOutBytes_{0};
  unsigned int frameCnt_{0};
  TDiff tdiffUncompress_;
  TDiff tdiffEncode_;
  TDiff tdiffDebayer_;
  TDiff tdiffFrameCopy_;
  TDiff tdiffSendFrame_;
  TDiff tdiffReceivePacket_;
  TDiff tdiffCopyOut_;
  TDiff tdiffPublish_;
  TDiff tdiffTotal_;
};
}  // namespace ffmpeg_encoder_decoder
#endif  // FFMPEG_ENCODER_DECODER__ENCODER_HPP_