@0biwank/screen-capture 2.0.1 → 2.0.2

package/src/HLSMuxer/VideoEncoder.cpp

@@ -1,462 +0,0 @@
-#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;
-}