VideoMuxerWithTasks.h

Go to the documentation of this file.

#pragma once
 
#include "AudioTools/AudioCodecs/ContainerCommon.h"
#include "AudioTools/Concurrency.h"
#include "AudioTools/Video/CodecVideo.h"
 
namespace audio_tools {

class VideoMuxerWithTasks {
 public:
  VideoMuxerWithTasks() = default;
  VideoMuxerWithTasks(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("VideoMuxerWithTasks: 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; }

  void setMutex(MutexBase &mutex) { p_mutex = &mutex; }

  void setTaskParameters(uint32_t stackSizeWords, uint8_t priority,
                         int core = -1) {
    task_stack_size = stackSizeWords;
    task_priority = priority;
    task_core = core;
  }

  bool begin() {
    if (p_muxer == nullptr) {
      LOGE("VideoMuxerWithTasks: setMuxer() was not called");
      return false;
    }
 
    bool have_video = false;
    if (p_video != nullptr) {
      VideoInfo info = p_video->videoInfo();
      if (info.fps <= 0 || info.width == 0 || info.height == 0) {
        LOGE(
            "VideoMuxerWithTasks: 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;
      // with a video encoder, the source's own format describes the raw
      // picture, not the container - query the encoder's own
      // VideoEncoder::videoInfo() instead, the reliable way to determine
      // what actually ends up in the container
      cfg.format = p_video_encoder != nullptr
                       ? p_video_encoder->videoInfo().format
                       : info.format;
      p_muxer->setVideoInfo(cfg);
      video_interval_ms = (uint32_t)(1000.0f / info.fps);
      video_sink.begin(*p_muxer);
      have_video = true;
    }
 
    bool 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);
    }
 
    if (!p_muxer->begin()) {
      LOGE("VideoMuxerWithTasks: Muxer::begin() failed");
      return false;
    }
 
    bool ok = true;
    if (have_video) {
      if (p_video_encoder != nullptr) {
        // routed through video_sink, not *p_muxer directly, so it can
        // determine each frame's real isKeyFrame status (see
        // MuxerVideoSink) instead of defaulting to true
        p_video_encoder->setOutput(video_sink);
        p_video_encoder->begin();
      }
      ok &= video_task.create("VideoMuxerVideoTask", task_stack_size,
                              task_priority, task_core);
      ok &= video_task.begin([this]() { videoTaskLoop(); });
    }
    if (have_audio) {
      if (p_audio_encoder != nullptr) {
        p_audio_encoder->setAudioInfo(audio_info);
        audio_sink.begin(*p_muxer);
        p_audio_encoder->setOutput(audio_sink);
        p_audio_encoder->begin();
      }
      ok &= audio_task.create("VideoMuxerAudioTask", task_stack_size,
                              task_priority, task_core);
      ok &= audio_task.begin([this]() { audioTaskLoop(); });
    }
    return ok;
  }

  void end() {
    audio_task.end();
    video_task.end();
    if (p_muxer != nullptr) p_muxer->end();
  }
 
  bool isRunning() { return audio_task.isRunning() || video_task.isRunning(); }
 
 protected:
  static constexpr size_t kAudioReadChunkSize = 512;

#ifdef __linux__
  using DefaultMutex = StdMutex;
#else
  using DefaultMutex = Mutex;
#endif
 
  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;
  MuxerAudioSink audio_sink;
  DefaultMutex default_mutex;
  MutexBase *p_mutex = &default_mutex;
  Task audio_task;
  Task video_task;
  uint32_t task_stack_size = 4096;
  uint8_t task_priority = 2;
  int task_core = -1;
  uint32_t video_interval_ms = 33;

  void audioTaskLoop() {
    if (p_audio != nullptr) {
      uint8_t buf[kAudioReadChunkSize];
      size_t n = p_audio->readBytes(buf, sizeof(buf));
      if (n > 0) {
        LockGuard guard(p_mutex);
        p_muxer->setStreamType(StreamContentType::Audio);
        if (p_audio_encoder != nullptr) {
          p_audio_encoder->write(buf, n);
        } else {
          p_muxer->addAudioFrame(buf, n);
        }
      } else {
        delay(1);
      }
    }
  }

  void videoTaskLoop() {
    uint32_t start = millis();
    if (p_video != nullptr && p_muxer != nullptr) {
      size_t len = 0;
      const uint8_t *data = p_video->nextFrame(len);
      if (data != nullptr && len > 0) {
        LockGuard guard(p_mutex);
        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);
        }
      }
    }
    uint32_t elapsed = millis() - start;
    // always yield, even when encoding took longer than the frame interval:
    // otherwise lower priority tasks (e.g. the idle task watched by the
    // task watchdog) never get to run on this core
    delay(elapsed < video_interval_ms ? video_interval_ms - elapsed : 1);
  }
};
 
}  // namespace audio_tools