rtsp_capture.h

#ifndef RTSP_CAPTURE_H
#define RTSP_CAPTURE_H

#pragma once
#include <string>
#include <thread>
#include <atomic>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
}

class RtspCapture {
 public:
  RtspCapture(const std::string& url);
  ~RtspCapture();

  void take(uint32_t width, uint32_t height);

 private:
  bool initFFmpeg();
  void cleanup();
  bool saveFrame(const std::string& filename, AVFrame* frame, uint32_t width, uint32_t height);

 private:
  std::string m_url;
  std::string m_outputDir;

  AVFormatContext* m_formatContext;
  AVCodecContext* m_codecContext;
  SwsContext* m_swsContext;
  int m_videoStreamIndex;
  std::string m_rtspUrl;
};

#endif // RTSPCAPTURE_H

rtsp_capture.cpp

#include "rtsp_capture.h"
#include <chrono>
#include <iostream>

RtspCapture::RtspCapture(const std::string& url)
    : m_url(url)
      , m_formatContext(nullptr)
      , m_codecContext(nullptr)
      , m_swsContext(nullptr)
      , m_videoStreamIndex(-1)
{
  initFFmpeg();
}

RtspCapture::~RtspCapture() {
  cleanup();
}

bool RtspCapture::initFFmpeg() {

#if LIBAVFORMAT_VERSION_MAJOR < 58
  av_register_all(); // 对于旧版本，显式调用以确保所有编解码器和格式被注册
#endif

  avformat_network_init();

  // 设置RTSP超时选项
  AVDictionary* options = nullptr;
  // av_dict_set(&options, "rtsp_transport", "tcp", 0);
  av_dict_set(&options, "stimeout", "30000000", 0);
  av_dict_set(&options, "max_delay", "500000", 0);  // 设置最大延迟
  av_dict_set(&options, "buffer_size", "1024000", 0);  // 设置缓冲区大小

  // av_log_set_level(AV_LOG_DEBUG);

  // 打开RTSP流
  m_formatContext = avformat_alloc_context();
  if (avformat_open_input(&m_formatContext, m_url.c_str(), nullptr, &options) != 0) {
    std::cerr << "无法打开RTSP流:" << m_url << std::endl;
    return false;
  }

  std::cout << "打开RTSP流成功:" << m_url << std::endl;

  // 获取流信息
  if (avformat_find_stream_info(m_formatContext, nullptr) < 0) {
    std::cerr << "无法获取流信息" << std::endl;
    return false;
  }

  // 查找视频流
  m_videoStreamIndex = -1;
  for (unsigned int i = 0; i < m_formatContext->nb_streams; i++) {
    if (m_formatContext->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
      m_videoStreamIndex = i;
      break;
    }
  }

  if (m_videoStreamIndex == -1) {
    std::cerr << "找不到视频流" << std::endl;
    return false;
  }

  // 获取解码器
  const AVCodec* codec = avcodec_find_decoder(m_formatContext->streams[m_videoStreamIndex]->codecpar->codec_id);
  if (!codec) {
    std::cerr << "找不到解码器" << std::endl;
    return false;
  }

  // 创建解码器上下文
  m_codecContext = avcodec_alloc_context3(codec);
  avcodec_parameters_to_context(m_codecContext, m_formatContext->streams[m_videoStreamIndex]->codecpar);

  // 打开解码器
  if (avcodec_open2(m_codecContext, codec, nullptr) < 0) {
    std::cerr << "无法打开解码器" << std::endl;
    return false;
  }

  std::cout << "FFmpeg初始化成功!" << std::endl;

  return true;
}

static uint64_t GetCurTick()
{
  auto now = std::chrono::system_clock::now();
  auto duration = now.time_since_epoch();
  return std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
}

void RtspCapture::take(uint32_t width, uint32_t height) {
  AVPacket* packet = av_packet_alloc();
  AVFrame* frame = av_frame_alloc();
  if (av_read_frame(m_formatContext, packet) >= 0) {
    if (packet->stream_index == m_videoStreamIndex) {
      // 解码视频帧
      if (avcodec_send_packet(m_codecContext, packet) >= 0) {
        if (avcodec_receive_frame(m_codecContext, frame) >= 0) {
          // 生成文件名
          std::string filename = "img.jpg";
          // std::string filename = "snapshot_" +
          //                        std::to_string(GetCurTick()) + ".jpg";

          // 保存帧为JPEG
          saveFrame(filename, frame, width, height);
        }
      }
    }
    av_packet_unref(packet);
  }

  av_frame_free(&frame);
  av_packet_free(&packet);

}

bool RtspCapture::saveFrame(const std::string& filename,  AVFrame* frame, uint32_t width, uint32_t height) {

  // 获取JPEG编码器
  const AVCodec* jpegCodec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
  if (!jpegCodec) {
    std::cerr << "无法找到JPEG编码器" << std::endl;
    return false;
  }

  // 创建编码器上下文
  AVCodecContext* jpegContext = avcodec_alloc_context3(jpegCodec);
  if (!jpegContext) {
    std::cerr << "无法创建JPEG编码器上下文" << std::endl;
    return false;
  }

  // 设置编码器参数
  jpegContext->width = width;
  jpegContext->height = height;
  jpegContext->time_base = (AVRational){1, 1};
  jpegContext->pix_fmt = AV_PIX_FMT_YUVJ420P;  // JPEG使用的颜色空间
  jpegContext->codec_type = AVMEDIA_TYPE_VIDEO;
  jpegContext->codec_id = AV_CODEC_ID_MJPEG;

  // 打开编码器
  if (avcodec_open2(jpegContext, jpegCodec, nullptr) < 0) {
    std::cerr << "无法打开JPEG编码器" << std::endl;
    avcodec_free_context(&jpegContext);
    return false;
  }

  // 分配转换后的帧
  AVFrame* yuvFrame = av_frame_alloc();
  yuvFrame->format = AV_PIX_FMT_YUVJ420P;
  yuvFrame->width = m_codecContext->width;
  yuvFrame->height = m_codecContext->height;
  av_frame_get_buffer(yuvFrame, 0);

  // 创建缩放上下文
  SwsContext* swsContext = sws_getContext(
      m_codecContext->width, m_codecContext->height, m_codecContext->pix_fmt,
      m_codecContext->width, m_codecContext->height, AV_PIX_FMT_YUVJ420P,
      SWS_BILINEAR, nullptr, nullptr, nullptr
      );

  if (!swsContext) {
    std::cerr << "无法创建转换上下文" << std::endl;
    av_frame_free(&yuvFrame);
    avcodec_free_context(&jpegContext);
    return false;
  }

  // 转换颜色空间
  sws_scale(swsContext, frame->data, frame->linesize,
            0, m_codecContext->height, yuvFrame->data, yuvFrame->linesize);

  // 编码为JPEG
  AVPacket* jpegPacket = av_packet_alloc();
  int ret = avcodec_send_frame(jpegContext, yuvFrame);
  if (ret < 0) {
    std::cerr << "发送帧到编码器失败" << std::endl;
    av_packet_free(&jpegPacket);
    av_frame_free(&yuvFrame);
    avcodec_free_context(&jpegContext);
    sws_freeContext(swsContext);
    return false;
  }

  ret = avcodec_receive_packet(jpegContext, jpegPacket);
  if (ret < 0) {
    std::cerr << "从编码器接收数据失败" << std::endl;
    av_packet_free(&jpegPacket);
    av_frame_free(&yuvFrame);
    avcodec_free_context(&jpegContext);
    sws_freeContext(swsContext);
    return false;
  }

  // 写入文件
  FILE* outFile = fopen(filename.c_str(), "wb");
  if (outFile) {
    fwrite(jpegPacket->data, 1, jpegPacket->size, outFile);
    fclose(outFile);
    // std::cout << "已保存图片: " << filename << std::endl;
  } else {
    std::cerr << "无法创建输出文件" << std::endl;
  }

  return true;
}

void RtspCapture::cleanup() {
  if (m_codecContext) {
    avcodec_free_context(&m_codecContext);
  }
  if (m_formatContext) {
    avformat_close_input(&m_formatContext);
    avformat_free_context(m_formatContext);
  }
}

main.cpp

#include "rtsp_capture.h"

int main()
{
    std::unique_ptr<RtspCapture> rtsp_capture_ptr_;
    std::string url = "rtsp://admin:[email protected]:554/live1";
    rtsp_capture_ptr_ = std::make_unique<RtspCapture>(url);
    rtsp_capture_ptr_->take(640, 480);

    getchar();
    return 0;
}