CodecH264.h

Go to the documentation of this file.

#pragma once
 
#include "AudioTools/CoreAudio/AudioBasic/Collections/Vector.h"
#include "AudioTools/Video/CodecVideo.h"
#include "TinyH264Decoder.h"
#include "TinyH264Encoder.h"
 
#ifndef H264_DEFAULT_ALLOCATOR
#ifdef ESP32
#define H264_DEFAULT_ALLOCATOR tinyh264::PSRAMAllocatorESP32<uint8_t>
#else
#define H264_DEFAULT_ALLOCATOR tinyh264::StdAllocator<uint8_t>
#endif
#endif
 
namespace audio_tools {
 
class H264Decoder : public VideoDecoder, public VideoInfoSource {
 public:
  H264Decoder() { decoder_.setCallback(onFrame, this); }
 
  H264Decoder(Print &out) : H264Decoder() { setOutput(out); }
 
  H264Decoder(VideoOutput &out) : H264Decoder() { setOutput(out); }
 
  VideoFormat codecFormat() override { return VideoFormat::H264; }
 
  bool isValid(const uint8_t *data, size_t len) override {
    for (size_t i = 0; i + 3 < len; i++) {
      if (data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 1) {
        uint8_t header = data[i + 3];
        if ((header & 0x80) == 0) {  // forbidden_zero_bit must be 0
          uint8_t nal_type = header & 0x1F;
          if (nal_type >= 1 && nal_type <= 23) return true;
        }
        i += 2;
      }
    }
    return false;
  }
 
  void setOutput(Print &out) override { p_out = &out; }
  void setOutput(VideoOutput &out) override { p_out_video = &out; }
 
 
  void setVideoFormat(VideoFormat format) override {
    switch (format) {
      case VideoFormat::RGB565:
      case VideoFormat::RGB666:
      case VideoFormat::RGB888:
      case VideoFormat::I420:
        pixel_format = format;
        break;
      default:
        LOGW("H264Decoder: unsupported VideoFormat %d", (int)format);
        break;
    }
  }
 
  VideoInfo videoInfo() override {
    VideoInfo info;
    info.format = pixel_format;
    info.width = (uint16_t)decoder_.width();
    info.height = (uint16_t)decoder_.height();
    return info;
  }
 
  void setMaxRefFrames(int n) { decoder_.setMaxRefFrames(n); }
 
  bool begin() override {
    decoder_.begin();
    return true;
  }
 
  void end() override { decoder_.end(); }
 
  size_t write(const uint8_t *data, size_t len) override {
    uint32_t decodeStart = millis();
    uint32_t outputMsBefore = total_output_ms_;
    decoder_.write(data, len);
    uint32_t outputMsDuringThisCall = total_output_ms_ - outputMsBefore;
    total_decode_ms_ += (millis() - decodeStart) - outputMsDuringThisCall;
    return len;
  }
 
  bool hasError() { return decoder_.hasError(); }
 
  bool isKeyFrame(const uint8_t *data, size_t len) override {
    return isH264KeyFrame(data, len);
  }
 
  uint64_t totalDecodeMs() const override { return total_decode_ms_; }
 
  tinyh264::TinyH264Decoder<H264_DEFAULT_ALLOCATOR> &driver() { return decoder_; }
 
 protected:
  uint64_t total_decode_ms_ = 0;
  // Accumulated time spent inside writeToOutput() (pixel conversion +
  // push to the configured output) since begin() - see totalDecodeMs()'s
  // own comment. write() reads the delta across its own call and
  // subtracts it, so total_decode_ms_ stays decode-only regardless of
  // how expensive the configured output's own write() turns out to be.
  uint32_t total_output_ms_ = 0;
 
  tinyh264::TinyH264Decoder<H264_DEFAULT_ALLOCATOR> decoder_;
  Print *p_out = nullptr;
  VideoOutput *p_out_video = nullptr;
  VideoFormat pixel_format = VideoFormat::RGB565;
  Vector<uint8_t> frame_buffer;
 
  static void onFrame(tinyh264::TinyH264Decoder<H264_DEFAULT_ALLOCATOR> &decoder,
                       void *userData) {
    static_cast<H264Decoder *>(userData)->writeToOutput();
  }
 
  void writeToOutput() {
    if (p_out == nullptr && p_out_video == nullptr) return;
    // Every decoded picture is always shown - see PacedVideoOutput
    // (Video/PacedVideoOutput.h) for how pacing/scheduling actually
    // happens.
    // Timed separately from decode (see total_output_ms_'s own comment) -
    // covers both the pixel-format conversion below and the push to
    // setOutput()'s target, i.e. render, if that target is a real
    // display.
    uint32_t outputStart = millis();
    size_t w = decoder_.width();
    size_t h = decoder_.height();
    size_t n = 0;
    switch (pixel_format) {
      case VideoFormat::RGB565: {
        size_t needed = w * h;
        if (frame_buffer.size() < needed * 2) {
          frame_buffer.resize(needed * 2);
          if (frame_buffer.data() == nullptr) {
            LOGE("H264Decoder: frame buffer allocation failed (RGB565)");
            break;
          }
        }
        n = decoder_.toRGB565((uint16_t *)frame_buffer.data(), needed);
        if (n > 0) writeToOutput(frame_buffer.data(), n * 2);
        break;
      }
      case VideoFormat::RGB666: {
        size_t needed = w * h * 3;
        if (frame_buffer.size() < needed) {
          frame_buffer.resize(needed);
          if (frame_buffer.data() == nullptr) {
            LOGE("H264Decoder: frame buffer allocation failed (RGB666)");
            break;
          }
        }
        n = decoder_.toRGB666(frame_buffer.data(), needed);
        if (n > 0) writeToOutput(frame_buffer.data(), n);
        break;
      }
      case VideoFormat::RGB888: {
        size_t needed = w * h * 3;
        if (frame_buffer.size() < needed) {
          frame_buffer.resize(needed);
          if (frame_buffer.data() == nullptr) {
            LOGE("H264Decoder: frame buffer allocation failed (RGB888)");
            break;
          }
        }
        n = decoder_.toRGB888(frame_buffer.data(), needed);
        if (n > 0) writeToOutput(frame_buffer.data(), n);
        break;
      }
      case VideoFormat::I420: {
        size_t needed = w * h + 2 * (w / 2) * (h / 2);
        if (frame_buffer.size() < needed) {
          frame_buffer.resize(needed);
          if (frame_buffer.data() == nullptr) {
            LOGE("H264Decoder: frame buffer allocation failed (I420)");
            break;
          }
        }
        n = decoder_.toYUV420(frame_buffer.data(), needed);
        if (n > 0) writeToOutput(frame_buffer.data(), n);
        break;
      }
      default:
        break;
    }
    total_output_ms_ += millis() - outputStart;
  }
 
  size_t writeToOutput(uint8_t *data, size_t len) {
    size_t result = 0;
    if (p_out != nullptr) result = p_out->write((const uint8_t *)data, len);
    if (p_out_video != nullptr) result += p_out_video->write((const uint8_t *)data, len);
    return result;
  }
};
 
class H264Encoder : public VideoEncoder {
 public:
  static constexpr size_t kDefaultBitstreamBufferSize = 16384;
 
  H264Encoder(int width = 0, int height = 0, int keyframeInterval = 0)
      : encoder_(width, height, keyframeInterval) {
    width_ = width;
    height_ = height;
  }
 
  H264Encoder(Print &out) {
    setOutput(out);
  }
 
  H264Encoder(VideoOutput &out) {
    setOutput(out);
  }
 
  void setOutput(Print &out) override { p_out = &out; }
  void setOutput(VideoOutput &out)  { p_out_video = &out; }
 
  void setVideoFormat(VideoFormat format) override {
    switch (format) {
      case VideoFormat::I420:
      case VideoFormat::YUV422:
      case VideoFormat::RGB565:
      case VideoFormat::RGB666:
      case VideoFormat::RGB888:
        input_format = format;
        break;
      default:
        LOGW("H264Encoder: unsupported VideoFormat %d", (int)format);
        break;
    }
  }
 
  VideoInfo videoInfo() override {
    VideoInfo info;
    info.format = VideoFormat::H264;
    info.width = (uint16_t)width_;
    info.height = (uint16_t)height_;
    return info;
  }
 
  void setSize(int width, int height) {
    width_ = width;
    height_ = height;
    encoder_.setSize(width, height);
  }
  void setQp(int qp) { encoder_.setQp(qp); }
  void setTargetBitrate(int bitsPerSecond, double fps) {
    encoder_.setTargetBitrate(bitsPerSecond, fps);
  }
  void setKeyframeInterval(int frames) {
    encoder_.setKeyframeInterval(frames);
  }
  void setStride(int strideY, int strideC) {
    encoder_.setStride(strideY, strideC);
  }
  void setPackedStride(int stride) { encoder_.setPackedStride(stride); }
 
  void setBitstreamBufferSize(size_t size) {
    if (size > bitstream.size()) {
      bitstream.resize(size);
      if (bitstream.data() == nullptr) {
        LOGE("H264Encoder: bitstream buffer allocation failed");
      }
    }
  }
 
  bool begin() override { return begin(false); }
  bool begin(bool reserveColorConversionScratch) {
    encoder_.begin(reserveColorConversionScratch);
    return true;
  }
  void end() override { encoder_.end(); }
 
  int lastQp() const { return encoder_.lastQp(); }
 
  size_t write(const uint8_t *data, size_t len) override {
    switch (input_format) {
      case VideoFormat::I420: {
        size_t ySize = (size_t)width_ * height_;
        size_t cSize = (size_t)(width_ / 2) * (height_ / 2);
        if (len < ySize + 2 * cSize) return 0;
        const uint8_t *srcY = data;
        const uint8_t *srcU = data + ySize;
        const uint8_t *srcV = data + ySize + cSize;
        return writeOut([&](uint8_t *dst, size_t cap) {
          return encoder_.encodeFrame(srcY, srcU, srcV, dst, cap);
        });
      }
      case VideoFormat::YUV422:
        return writeOut([&](uint8_t *dst, size_t cap) {
          return encoder_.encodeFrameYuv422(data, dst, cap);
        });
      case VideoFormat::RGB565:
        return writeOut([&](uint8_t *dst, size_t cap) {
          return encoder_.encodeFrameRgb565((const uint16_t *)data, dst, cap);
        });
      case VideoFormat::RGB666:
        return writeOut([&](uint8_t *dst, size_t cap) {
          return encoder_.encodeFrameRgb666(data, dst, cap);
        });
      case VideoFormat::RGB888:
        return writeOut([&](uint8_t *dst, size_t cap) {
          return encoder_.encodeFrameRgb888(data, dst, cap);
        });
      default:
        return 0;
    }
  }
 
  tinyh264::TinyH264Encoder<H264_DEFAULT_ALLOCATOR> &driver() { return encoder_; }
 
 protected:
  static constexpr int kMaxGrowRetries = 4;
 
  tinyh264::TinyH264Encoder<H264_DEFAULT_ALLOCATOR> encoder_;
  Print *p_out = nullptr;
  VideoOutput *p_out_video = nullptr;
  Vector<uint8_t> bitstream;
  VideoFormat input_format = VideoFormat::I420;
  int width_ = 0;
  int height_ = 0;
 
  template <typename Fn>
  size_t writeOut(Fn encodeInto) {
    if (p_out == nullptr && p_out_video == nullptr) return 0;
    if (bitstream.size() == 0) {
      bitstream.resize(kDefaultBitstreamBufferSize);
      if (bitstream.data() == nullptr) {
        LOGE("H264Encoder: bitstream buffer allocation failed");
        return 0;
      }
    }
    size_t n = encodeInto(bitstream.data(), bitstream.size());
    for (int attempt = 0; n == 0 && attempt < kMaxGrowRetries; ++attempt) {
      bitstream.resize(bitstream.size() * 2);
      if (bitstream.data() == nullptr) {
        LOGE("H264Encoder: bitstream buffer grow failed");
        return 0;
      }
      n = encodeInto(bitstream.data(), bitstream.size());
    }
    if (p_out !=nullptr && n > 0) p_out->write(bitstream.data(), n);
    if (p_out_video !=nullptr && n > 0) p_out_video->write(bitstream.data(), n);
    return n;
  }
};
 
}  // namespace audio_tools