@0biwank/screen-capture 1.0.1 → 2.0.1

package/src/HLSMuxer/VideoEncoder.cpp

@@ -0,0 +1,462 @@
+#include "HLSMuxer/VideoEncoder.h"
+#include <iostream>
+#include <sstream>
+#include <iomanip>
+#include <string>
+
+extern "C"
+{
+#include <libavcodec/avcodec.h>
+#include <libavformat/avformat.h>
+#include <libswscale/swscale.h>
+#include <libavutil/opt.h>
+#include <libavutil/imgutils.h>
+#include <libavutil/frame.h>
+}
+
+VideoEncoder::VideoEncoder()
+{
+  m_codecContext = nullptr;
+  m_codec = nullptr;
+  m_frame = nullptr;
+  m_swsContext = nullptr;
+  m_formatContext = nullptr;
+  m_stream = nullptr;
+
+  m_frameCount = 0;
+  m_pts = 0;
+  m_dts = 0;
+  m_lastInputPts = -1;
+  currentChunkIndex = 0;
+
+  m_initialized = false;
+  isPaused = false;
+}
+
+VideoEncoder::~VideoEncoder()
+{
+  cleanup();
+  while (!m_packetQueue.empty())
+  {
+    AVPacket *pkt = m_packetQueue.front();
+    m_packetQueue.pop();
+    av_packet_free(&pkt);
+  }
+}
+
+bool VideoEncoder::initialize(int width, int height, int fps, int bitrate,
+                              OutputMode outputMode, int gopSize)
+{
+  m_codec = avcodec_find_encoder(AV_CODEC_ID_H264);
+  if (!m_codec)
+  {
+    std::cerr << "H264 encoder not found!" << std::endl;
+    return false;
+  }
+  m_codecContext = avcodec_alloc_context3(m_codec);
+  if (!m_codecContext)
+  {
+    std::cerr << "Could not allocate codec context" << std::endl;
+    return false;
+  }
+  m_codecContext->bit_rate = bitrate;
+  m_codecContext->width = width;
+  m_codecContext->height = height;
+  m_codecContext->time_base = {1, fps};
+  m_codecContext->framerate = {fps, 1};
+  m_codecContext->gop_size = gopSize > 0 ? gopSize : fps;
+  m_codecContext->max_b_frames = 0;
+  m_codecContext->pix_fmt = AV_PIX_FMT_YUV420P;
+  this->frameWidth = width;
+  this->frameHeight = height;
+  this->frameRate = fps;
+  this->mode = outputMode;
+
+  AVDictionary *codecOptions = nullptr;
+  if (m_codec && m_codec->name && std::string(m_codec->name) == "libx264")
+  {
+    av_dict_set(&codecOptions, "preset", "veryfast", 0);
+    av_dict_set(&codecOptions, "tune", "zerolatency", 0);
+  }
+
+  int ret = avcodec_open2(m_codecContext, m_codec, &codecOptions);
+  av_dict_free(&codecOptions);
+  if (ret < 0)
+  {
+    std::cerr << "Could not open codec" << std::endl;
+    cleanup();
+    return false;
+  }
+  m_swsContext = sws_getContext(width, height, AV_PIX_FMT_BGRA, width, height,
+                                AV_PIX_FMT_YUV420P, SWS_BILINEAR, nullptr,
+                                nullptr, nullptr);
+  if (!m_swsContext)
+  {
+    std::cerr << "Could not initialize sws context" << std::endl;
+    cleanup();
+    return false;
+  }
+  m_frame = av_frame_alloc();
+  if (!m_frame)
+  {
+    cleanup();
+    return false;
+  }
+  m_frame->format = m_codecContext->pix_fmt;
+  m_frame->width = width;
+  m_frame->height = height;
+
+  ret = av_frame_get_buffer(m_frame, 0);
+  if (ret < 0)
+  {
+    std::cerr << "Could not allocate frame buffer" << std::endl;
+    cleanup();
+    return false;
+  }
+
+  if (outputMode == OutputMode::STANDALONE)
+  {
+    // Initialize format context for TS output
+    avformat_alloc_output_context2(&m_formatContext, nullptr, "mpegts", nullptr);
+    if (!m_formatContext)
+    {
+      std::cerr << "Could not create output context for TS" << std::endl;
+      cleanup();
+      return false;
+    }
+
+    // Create video stream in the TS container
+    m_stream = avformat_new_stream(m_formatContext, nullptr);
+    if (!m_stream)
+    {
+      std::cerr << "Could not create video stream" << std::endl;
+      cleanup();
+      return false;
+    }
+
+    // Copy codec parameters to stream
+    avcodec_parameters_from_context(m_stream->codecpar, m_codecContext);
+    m_stream->time_base = m_codecContext->time_base;
+
+    targetChunkDuration = 10.0; // 10 second chunks by default
+    currentSegmentStartPts = 0;
+    currentSegmentDuration = 0.0;
+  }
+
+  m_initialized = true;
+  return true;
+}
+
+bool VideoEncoder::encodeFrame(uint8_t *pixelData, int stride, int64_t pts)
+{
+  if (!m_initialized)
+  {
+    std::cerr << "Encoder not initalized!" << std::endl;
+    return false;
+  }
+
+  int ret = av_frame_make_writable(m_frame);
+  if (ret < 0)
+  {
+    std::cerr << "Frame not writeable" << std::endl;
+    return false;
+  }
+
+  const uint8_t *srcData[4] = {pixelData, nullptr, nullptr, nullptr};
+  int srcLinesize[4] = {stride, 0, 0, 0};
+  sws_scale(
+      m_swsContext,
+      srcData,
+      srcLinesize,
+      0,
+      frameHeight,
+      m_frame->data,
+      m_frame->linesize);
+
+  if (pts == -1)
+  {
+    m_frame->pts = m_frameCount;
+  }
+  else
+  {
+    m_frame->pts = pts;
+  }
+
+  if (m_frame->pts <= m_lastInputPts)
+  {
+    m_frame->pts = m_lastInputPts + 1;
+  }
+  m_lastInputPts = m_frame->pts;
+
+  m_pts = m_frame->pts;
+
+  ret = avcodec_send_frame(m_codecContext, m_frame);
+  // Reset forced keyframe flag after submission so only one frame gets the IDR hint.
+  m_frame->pict_type = AV_PICTURE_TYPE_NONE;
+  m_frame->flags &= ~AV_FRAME_FLAG_KEY;
+  if (ret < 0)
+  {
+    std::cerr << "Error sending frame to encode" << std::endl;
+    return false;
+  }
+
+  while (ret >= 0)
+  {
+    AVPacket *packet = av_packet_alloc();
+    ret = avcodec_receive_packet(m_codecContext, packet);
+    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
+    {
+      av_packet_free(&packet);
+      break;
+    }
+
+    if (mode == OutputMode::MUXED)
+    {
+      m_packetQueue.push(packet);
+    }
+    else if (mode == OutputMode::STANDALONE)
+    {
+      // Check if we need to start a new chunk
+      if (shouldStartNewChunk() || currentChunkIndex == 0)
+      {
+        // Close current chunk if exists
+        if (m_formatContext && m_formatContext->pb)
+        {
+          av_write_trailer(m_formatContext);
+          avio_closep(&m_formatContext->pb);
+
+          currentSegmentDuration = (m_pts - currentSegmentStartPts) * av_q2d(m_codecContext->time_base);
+        }
+
+        // Generate filename for new chunk
+        std::ostringstream filename;
+        filename << "video_chunk" << std::setfill('0') << std::setw(3) << currentChunkIndex << ".ts";
+        std::string chunkFilename = filename.str();
+        m_chunkFileNames.push_back(chunkFilename);
+        std::cout << "Creating video chunk: " << chunkFilename << std::endl;
+
+        // Open new chunk file
+        if (avio_open(&m_formatContext->pb, chunkFilename.c_str(), AVIO_FLAG_WRITE) < 0)
+        {
+          std::cerr << "Could not open output file: " << chunkFilename << std::endl;
+          av_packet_free(&packet);
+          break;
+        }
+
+        // Write header for new chunk
+        if (avformat_write_header(m_formatContext, nullptr) < 0)
+        {
+          std::cerr << "Could not write header for chunk: " << chunkFilename << std::endl;
+          av_packet_free(&packet);
+          break;
+        }
+
+        currentChunkIndex++;
+        currentSegmentStartPts = m_pts;
+
+        // Force keyframe at chunk boundaries
+        forceKeyFrame();
+      }
+
+      // Write packet to current chunk
+      packet->stream_index = m_stream->index;
+      av_packet_rescale_ts(packet, m_codecContext->time_base, m_stream->time_base);
+      av_interleaved_write_frame(m_formatContext, packet);
+      av_packet_free(&packet);
+    }
+  }
+  m_frameCount++;
+  return true;
+}
+
+void VideoEncoder::cleanup()
+{
+  std::cout << "VideoEncoder::cleanup() called, mode=" << (int)mode << ", initialized=" << m_initialized << std::endl;
+  
+  // Flush encoder to get remaining packets and finalize last segment
+  if (m_initialized && m_codecContext)
+  {
+    std::cout << "Flushing encoder..." << std::endl;
+    // Send NULL frame to flush encoder
+    int flushRet = avcodec_send_frame(m_codecContext, nullptr);
+    std::cout << "Flush send returned: " << flushRet << std::endl;
+    
+    // Receive all remaining packets
+    while (true)
+    {
+      AVPacket *packet = av_packet_alloc();
+      int ret = avcodec_receive_packet(m_codecContext, packet);
+      
+      if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
+      {
+        av_packet_free(&packet);
+        break;
+      }
+      
+      if (ret < 0)
+      {
+        av_packet_free(&packet);
+        break;
+      }
+      
+      if (mode == OutputMode::MUXED)
+      {
+        m_packetQueue.push(packet);
+      }
+      else if (mode == OutputMode::STANDALONE && m_formatContext)
+      {
+        // Open first chunk if not already open
+        if (!m_formatContext->pb && currentChunkIndex == 0)
+        {
+          std::ostringstream filename;
+          filename << "video_chunk" << std::setfill('0') << std::setw(3) << currentChunkIndex << ".ts";
+          std::string chunkFilename = filename.str();
+          m_chunkFileNames.push_back(chunkFilename);
+          std::cout << "Flush: Creating first video chunk: " << chunkFilename << std::endl;
+          
+          if (avio_open(&m_formatContext->pb, chunkFilename.c_str(), AVIO_FLAG_WRITE) < 0)
+          {
+            std::cerr << "Could not open output file: " << chunkFilename << std::endl;
+            av_packet_free(&packet);
+            continue;
+          }
+          
+          if (avformat_write_header(m_formatContext, nullptr) < 0)
+          {
+            std::cerr << "Flush: Could not write header for output file" << std::endl;
+            av_packet_free(&packet);
+            continue;
+          }
+
+          currentChunkIndex++;
+          currentSegmentStartPts = m_pts;
+        }
+        
+        if (m_formatContext->pb)
+        {
+          std::cout << "Flush: Writing packet to STANDALONE" << std::endl;
+          packet->stream_index = m_stream->index;
+          av_packet_rescale_ts(packet, m_codecContext->time_base, m_stream->time_base);
+          av_interleaved_write_frame(m_formatContext, packet);
+        }
+        av_packet_free(&packet);
+      }
+      else
+      {
+        av_packet_free(&packet);
+      }
+    }
+    
+    // Finalize last segment if in STANDALONE mode
+    if (mode == OutputMode::STANDALONE && m_formatContext && m_formatContext->pb)
+    {
+      // Update duration for the final segment
+      currentSegmentDuration = (m_pts - currentSegmentStartPts) * av_q2d(m_codecContext->time_base);
+      
+      av_write_trailer(m_formatContext);
+      avio_closep(&m_formatContext->pb);
+    }
+  }
+  
+  if (mode != OutputMode::MUXED)
+  {
+    while (!m_packetQueue.empty())
+    {
+      AVPacket *pkt = m_packetQueue.front();
+      m_packetQueue.pop();
+      av_packet_free(&pkt);
+    }
+  }
+  if (m_frame)
+  {
+    av_frame_free(&m_frame);
+    m_frame = nullptr;
+  }
+  if (m_swsContext)
+  {
+    sws_freeContext(m_swsContext);
+    m_swsContext = nullptr;
+  }
+  if (m_codecContext)
+  {
+    avcodec_free_context(&m_codecContext);
+    m_codecContext = nullptr;
+  }
+  if (m_formatContext)
+  {
+    if (m_formatContext->pb)
+    {
+      // Should already be closed above, but ensure it's closed
+      avio_closep(&m_formatContext->pb);
+    }
+    avformat_free_context(m_formatContext);
+    m_formatContext = nullptr;
+  }
+  m_initialized = false;
+}
+
+bool VideoEncoder::hasPackets() const
+{
+  return !m_packetQueue.empty();
+}
+
+AVPacket *VideoEncoder::getNextPacket()
+{
+  if (m_packetQueue.empty())
+  {
+    return nullptr;
+  }
+
+  AVPacket *packet = m_packetQueue.front();
+  m_packetQueue.pop();
+  return packet;
+}
+
+void VideoEncoder::forceKeyFrame()
+{
+  if (m_frame && m_codecContext)
+  {
+    m_frame->pict_type = AV_PICTURE_TYPE_I;
+    m_frame->flags |= AV_FRAME_FLAG_KEY;
+  }
+}
+
+bool VideoEncoder::shouldStartNewChunk() const
+{
+  double duration = (m_pts - currentSegmentStartPts) * av_q2d(m_codecContext->time_base);
+  return duration >= targetChunkDuration;
+}
+
+std::string VideoEncoder::generatePlaylist() const
+{
+  if (mode != OutputMode::STANDALONE || m_chunkFileNames.empty())
+  {
+    return "";
+  }
+
+  std::string playlist = "#EXTM3U\n";
+  playlist += "#EXT-X-VERSION:3\n";
+  playlist += "#EXT-X-TARGETDURATION:" + std::to_string(static_cast<int>(targetChunkDuration + 1)) + "\n";
+  playlist += "#EXT-X-MEDIA-SEQUENCE:0\n";
+
+  for (size_t i = 0; i < m_chunkFileNames.size(); ++i)
+  {
+    // Calculate actual duration for each segment
+    double duration = targetChunkDuration; // Default duration
+    if (i == m_chunkFileNames.size() - 1 && currentSegmentDuration > 0)
+    {
+      duration = currentSegmentDuration; // Last segment might be shorter
+    }
+
+    playlist += "#EXTINF:" + std::to_string(duration) + ",\n";
+    playlist += m_chunkFileNames[i] + "\n";
+  }
+
+  // Add end list marker if encoding is complete
+  if (!m_initialized)
+  {
+    playlist += "#EXT-X-ENDLIST\n";
+  }
+
+  return playlist;
+}