VideoMuxer.h

Go to the documentation of this file.

#pragma once
 
#include "AudioTools/AudioCodecs/ContainerCommon.h"
#include "AudioTools/Video/CodecVideo.h"
 
namespace audio_tools {
 
class VideoMuxer {
 public:
  VideoMuxer() = default;
  VideoMuxer(Muxer& muxer) { setMuxer(muxer); }
 
  void setMuxer(Muxer& muxer) { p_muxer = &muxer; }
 
  void setAudioSource(Stream& pcm) { p_audio = &pcm; }
 
  void setAudioSource(AudioStream& out) {
    p_audio = &out;
    p_source_audio_info = &out;
    audio_info = out.audioInfo();
  }
 
  void setAudioInfo(AudioInfo info) {
    audio_info = info;
    if (p_source_audio_info) {
      p_source_audio_info->setAudioInfo(info);
    } else {
      LOGI("VideoMuxer: setAudioInfo() source not updated");
    }
  }
 
  void setAudioEncoder(AudioEncoder& encoder) { p_audio_encoder = &encoder; }
 
  void setVideoSource(VideoFrameSource& video) { p_video = &video; }
 
  void setVideoEncoder(VideoEncoder& encoder) { p_video_encoder = &encoder; }
 
  bool begin() {
    if (p_muxer == nullptr) {
      LOGE("VideoMuxer: setMuxer() was not called");
      return false;
    }
 
    have_video = false;
    if (p_video != nullptr) {
      VideoInfo info = p_video->videoInfo();
      if (info.fps <= 0 || info.width == 0 || info.height == 0) {
        LOGE(
            "VideoMuxer: VideoFrameSource::videoInfo() incomplete "
            "(width=%d height=%d fps=%f)",
            (int)info.width, (int)info.height, info.fps);
        return false;
      }
      MuxerVideoConfig cfg;
      cfg.width = info.width;
      cfg.height = info.height;
      cfg.fps = info.fps;
      cfg.format = p_video_encoder != nullptr
                       ? p_video_encoder->videoInfo().format
                       : info.format;
      p_muxer->setVideoInfo(cfg);
      video_fps = info.fps;
      video_interval_ms = (uint32_t)(1000.0f / info.fps);
      video_sink.begin(*p_muxer);
      have_video = true;
    }
 
    have_audio = p_audio != nullptr;
    if (have_audio) {
      AudioInfoFormat muxer_audio_info;
      if (p_audio_encoder != nullptr) {
        muxer_audio_info = AudioInfoFormat(audio_info);
        AudioFormat mime_format = fromMime(p_audio_encoder->mime());
        muxer_audio_info.format = mime_format != AudioFormat::UNKNOWN
                                      ? mime_format
                                      : AudioFormat::PCM;
      } else {
        muxer_audio_info = AudioInfoFormat(audio_info);
        muxer_audio_info.format = AudioFormat::PCM;
      }
      p_muxer->setAudioInfo(muxer_audio_info);
      sample_rate = muxer_audio_info.sample_rate;
      bytes_per_sample =
          (muxer_audio_info.bits_per_sample / 8) * muxer_audio_info.channels;
    }
 
    if (!p_muxer->begin()) {
      LOGE("VideoMuxer: Muxer::begin() failed");
      return false;
    }
 
    if (have_video && p_video_encoder != nullptr) {
      // routed through video_sink, not *p_muxer directly, so it can
      // determine each frame's real isKeyFrame status (see
      // MuxerVideoSink, ContainerCommon.h) instead of defaulting to true
      p_video_encoder->setOutput(video_sink);
      p_video_encoder->begin();
    }
    if (have_audio && p_audio_encoder != nullptr) {
      p_audio_encoder->setAudioInfo(audio_info);
      p_audio_encoder->setOutput(*p_muxer);
      p_audio_encoder->begin();
    }
 
    audio_samples_written = 0;
    video_frame_count = 0;
    return true;
  }
 
  void end() {
    if (p_muxer != nullptr) p_muxer->end();
  }
 
  size_t copy() {
    if (have_audio) return copyAudio();
    if (have_video) return copyVideoPaced() ? 1 : 0;
    return 0;
  }
 
 protected:
  static constexpr size_t kAudioReadChunkSize = 512;
 
  Muxer* p_muxer = nullptr;
  Stream* p_audio = nullptr;
  AudioEncoder* p_audio_encoder = nullptr;
  AudioInfo audio_info;
  AudioInfoSupport* p_source_audio_info = nullptr;
  VideoFrameSource* p_video = nullptr;
  VideoEncoder* p_video_encoder = nullptr;
  MuxerVideoSink video_sink;
  bool have_audio = false;
  bool have_video = false;
  float video_fps = 0;
  uint32_t video_interval_ms = 33;
  uint32_t sample_rate = 0;
  uint32_t bytes_per_sample = 0;
  uint64_t audio_samples_written = 0;
  uint32_t video_frame_count = 0;
 
  size_t copyAudio() {
    uint8_t buf[kAudioReadChunkSize];
    size_t n = p_audio->readBytes(buf, sizeof(buf));
    if (n == 0) return 0;
 
    p_muxer->setStreamType(StreamContentType::Audio);
    if (p_audio_encoder != nullptr) {
      p_audio_encoder->write(buf, n);
    } else {
      p_muxer->addAudioFrame(buf, n);
    }
 
    if (bytes_per_sample > 0) audio_samples_written += n / bytes_per_sample;
 
    // catch up: mux/encode every video frame now due, given the audio
    // time just reached - this, not a timer, is what paces video. Strict
    // '<' (not '<=') so a duration landing exactly on a frame boundary
    // doesn't pull one extra frame for the instant at that boundary.
    if (have_video && sample_rate > 0) {
      double audio_time_s = (double)audio_samples_written / sample_rate;
      while ((double)video_frame_count / video_fps < audio_time_s) {
        if (!copyVideo()) break;  // source exhausted
      }
    }
    return n;
  }
 
  bool copyVideo() {
    size_t len = 0;
    const uint8_t* data = p_video->nextFrame(len);
    if (data == nullptr || len == 0) return false;
 
    p_muxer->setStreamType(StreamContentType::Video);
    if (p_video_encoder != nullptr) {
      p_video_encoder->write(data, len);
    } else {
      bool is_key = p_muxer->getVideoInfo().format == VideoFormat::H264
                        ? isH264KeyFrame(data, len)
                        : true;
      p_muxer->addVideoFrame(data, len, is_key);
    }
    video_frame_count++;
    return true;
  }
 
  bool copyVideoPaced() {
    uint32_t start = millis();
    if (!copyVideo()) return false;
    uint32_t elapsed = millis() - start;
    if (elapsed < video_interval_ms) delay(video_interval_ms - elapsed);
    return true;
  }
};
 
}  // namespace audio_tools