CodecMPG.h

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include "AudioTools/CoreAudio/AudioBasic/Collections/Vector.h"
#include "AudioTools/Video/CodecVideo.h"
 
#include "TinyMPGDecoder.h"
#include "TinyMPGEncoder.h"
 
 
#ifndef MPG_DEFAULT_ALLOCATOR
#ifdef ESP32
#  define MPG_DEFAULT_ALLOCATOR tinympg::PSRAMAllocatorESP32<uint8_t>
#else
#  define MPG_DEFAULT_ALLOCATOR tinympg::StdAllocator<uint8_t>
#endif
#endif
 
 
namespace audio_tools {
 
class MPGDecoder : public VideoDecoder, public VideoInfoSource {
 public:
  MPGDecoder() { decoder_.setCallback(onFrame, this); }
 
  MPGDecoder(Print &out) : MPGDecoder() { setOutput(out); }
 
  MPGDecoder(VideoOutput &out) : MPGDecoder() { setOutput(out); }
 
  VideoFormat codecFormat() override { return VideoFormat::MPEG1; }
 
  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 &&
          data[i + 3] == 0xB3) {
        return true;
      }
    }
    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("MPGDecoder: unsupported VideoFormat %d", (int)format);
        break;
    }
  }
 
  void setIgnorePFrames(bool active) { decoder_.setIgnorePFrames(active); }
  bool ignorePFrames() const { return decoder_.ignorePFrames(); }
 
  void setByteSwap(bool active) { byte_swap = active; }
  bool byteSwap() const { return byte_swap; }
 
  VideoInfo videoInfo() override {
    VideoInfo info;
    info.format = pixel_format;
    info.width = (uint16_t)decoder_.width();
    info.height = (uint16_t)decoder_.height();
    return info;
  }
 
  bool begin() override {
    decoder_.begin();
    return true;
  }
 
  void end() override {
    decoder_.write(nullptr, 0);
    decoder_.end();
  }
 
  size_t write(const uint8_t *data, size_t len) override {
    frame_emitted = false;
    decoder_.write(data, len);
    return len;
  }
 
  void flush() override {}
 
  bool hasError() { return decoder_.hasError(); }
 
  bool isKeyFrame(const uint8_t *data, size_t len) override {
    return isMpeg1KeyFrame(data, len);
  }
 
  bool hadOutput() const override { return frame_emitted; }
 
  tinympg::TinyMPGDecoder<MPG_DEFAULT_ALLOCATOR> &driver() { return decoder_; }
 
 protected:
  tinympg::TinyMPGDecoder<MPG_DEFAULT_ALLOCATOR> decoder_;
  Print *p_out = nullptr;
  VideoOutput *p_out_video = nullptr;
  VideoFormat pixel_format = VideoFormat::RGB565;
  bool byte_swap = true;  // see setByteSwap()
  bool frame_emitted = false;  // see hadOutput()
  Vector<uint8_t> frame_buffer;
 
  static void onFrame(tinympg::TinyMPGDecoder<MPG_DEFAULT_ALLOCATOR> &decoder,
                      void *userData) {
    static_cast<MPGDecoder *>(userData)->writeFrame(decoder);
  }
 
  void writeFrame(tinympg::TinyMPGDecoder<MPG_DEFAULT_ALLOCATOR> &decoder) {
    if (p_out == nullptr && p_out_video == nullptr) return;
 
    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("MPGDecoder: frame buffer allocation failed (RGB565)");
            return;
          }
        }
        n = decoder.toRGB565((uint16_t *)frame_buffer.data(), needed);
        if (n > 0) {
          if (byte_swap) {
            uint16_t *px = (uint16_t *)frame_buffer.data();
            for (size_t i = 0; i < n; i++) {
              uint16_t v = px[i];
              px[i] = (uint16_t)((v << 8) | (v >> 8));
            }
          }
          writeToOutput(frame_buffer.data(), n * 2);
        }
        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("MPGDecoder: frame buffer allocation failed (RGB888)");
            return;
          }
        }
        n = decoder.toRGB888(frame_buffer.data(), needed);
        if (n > 0) writeToOutput(frame_buffer.data(), n);
        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("MPGDecoder: frame buffer allocation failed (RGB666)");
            return;
          }
        }
        n = decoder.toRGB666(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("MPGDecoder: frame buffer allocation failed (I420)");
            return;
          }
        }
        n = decoder.toYUV420(frame_buffer.data(), needed);
        if (n > 0) writeToOutput(frame_buffer.data(), n);
        break;
      }
      default:
        break;
    }
    if (n > 0) frame_emitted = true;
  }
 
  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 MPGEncoder : public VideoEncoder {
 public:
  MPGEncoder(int width = 0, int height = 0) {
    width_ = width;
    height_ = height;
    encoder_.setQp(8);
    if (width_ > 0 && height_ > 0) {
      encoder_.setSize(width_, height_);
      ensureBitstreamBuffer();
    }
  }
 
  MPGEncoder(Print &out) { setOutput(out); }
 
  MPGEncoder(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("MPGEncoder: unsupported VideoFormat %d", (int)format);
        break;
    }
  }
 
  VideoInfo videoInfo() override {
    VideoInfo info;
    info.format = VideoFormat::MPEG1;
    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);
    ensureBitstreamBuffer();
  }
 
  void setQp(int qscale) { encoder_.setQp(qscale); }
 
  void setFrameRate(int fps) { encoder_.setFrameRate(fps); }
 
  void setBitstreamBufferSize(size_t bytes) {
    bitstream_buffer.resize(bytes);
    if (bitstream_buffer.data() == nullptr && bytes > 0) {
      LOGE("MPGEncoder: bitstream buffer allocation failed");
    }
  }
 
  bool begin() override {
    if (width_ <= 0 || height_ <= 0) {
      LOGE("MPGEncoder: invalid size - call setSize() first");
      return false;
    }
    if (!ensureBitstreamBuffer()){
      LOGE("MPGEncoder: failed to allocate bitstream buffer");
      return false;
    }
    return true;
  }
 
  void end() override {}
 
  size_t write(const uint8_t *data, size_t len) override {
    if (p_out == nullptr && p_out_video == nullptr) return 0;
    if (width_ <= 0 || height_ <= 0) {
      LOGE("MPGEncoder: invalid size - call setSize() first");
      return 0;
    }
    if (!ensureBitstreamBuffer()) {
      LOGE("MPGEncoder: bitstream buffer not available");
      return 0;
    }
 
    switch (input_format) {
      case VideoFormat::I420: {
        size_t y_size = (size_t)width_ * (size_t)height_;
        size_t uv_size = y_size / 4;
        size_t min_size = y_size + uv_size + uv_size;
        if (len < min_size) {
          LOGE("MPGEncoder: input buffer too small for I420 frame");
          return 0;
        }
        const uint8_t *srcY = data;
        const uint8_t *srcU = data + y_size;
        const uint8_t *srcV = data + y_size + uv_size;
        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:
        LOGE("MPGEncoder: unsupported VideoFormat %d", (int)input_format);
        return 0;
    }
  }
 
  tinympg::TinyMPGEncoder<MPG_DEFAULT_ALLOCATOR> &driver() { return encoder_; }
 
 protected:
  static constexpr int kMaxGrowRetries = 4;
 
  tinympg::TinyMPGEncoder<MPG_DEFAULT_ALLOCATOR> encoder_;
  Print *p_out = nullptr;
  VideoOutput *p_out_video = nullptr;
  VideoFormat input_format = VideoFormat::I420;
  int width_ = 0;
  int height_ = 0;
  Vector<uint8_t> bitstream_buffer;
 
  template <typename Fn>
  size_t writeOut(Fn encodeInto) {
    if (p_out == nullptr && p_out_video == nullptr) return 0;
    size_t encoded = encodeInto(bitstream_buffer.data(), bitstream_buffer.size());
    for (int attempt = 0; encoded == 0 && attempt < kMaxGrowRetries; ++attempt) {
      bitstream_buffer.resize(bitstream_buffer.size() * 2);
      if (bitstream_buffer.data() == nullptr) {
        LOGE("MPGEncoder: bitstream buffer grow failed");
        return 0;
      }
      encoded = encodeInto(bitstream_buffer.data(), bitstream_buffer.size());
    }
    if (encoded == 0) {
      LOGE("MPGEncoder: encode failed (status=%d)", (int)encoder_.lastStatus());
      return 0;
    }
    if (p_out != nullptr) p_out->write(bitstream_buffer.data(), encoded);
    if (p_out_video != nullptr) p_out_video->write(bitstream_buffer.data(), encoded);
    return encoded;
  }
 
  bool ensureBitstreamBuffer() {
    if (width_ <= 0 || height_ <= 0) return false;
    size_t needed = (size_t)width_ * (size_t)height_ * 2 + 4096;
    if (bitstream_buffer.size() < needed) {
      bitstream_buffer.resize(needed);
      if (bitstream_buffer.data() == nullptr) {
        return false;
      }
    }
    return bitstream_buffer.data() != nullptr;
  }
};
 
}  // namespace audio_tools