sladro
e13cb3659c
- 创建基本项目结构和目录 - 添加CMake构建系统 - 实现基础的配置解析功能 - 添加YOLO推理框架支持 - 集成RTSP和视频流处理功能 - 添加性能监控和日志系统
314 lines
8.1 KiB
C++
314 lines
8.1 KiB
C++
#include "rtsp_writer.hpp"
|
||
#include "../common/logger.hpp"
|
||
|
||
namespace pipeline {
|
||
|
||
RtspWriter::RtspWriter() = default;
|
||
|
||
RtspWriter::~RtspWriter() {
|
||
release();
|
||
}
|
||
|
||
bool RtspWriter::init(const OutputTargetConfig& config) {
|
||
std::lock_guard<std::mutex> lock(mutex_);
|
||
|
||
if (is_initialized_) {
|
||
Logger::warning("RtspWriter already initialized");
|
||
return true;
|
||
}
|
||
|
||
// 基本配置验证
|
||
if (config.name.empty()) {
|
||
setLastError("Empty name in config");
|
||
return false;
|
||
}
|
||
|
||
if (config.path.empty()) {
|
||
setLastError("Empty path in config");
|
||
return false;
|
||
}
|
||
|
||
if (config.fps <= 0) {
|
||
setLastError("Invalid fps in config: " + std::to_string(config.fps));
|
||
return false;
|
||
}
|
||
|
||
name_ = config.name;
|
||
url_ = config.path;
|
||
fps_ = config.fps;
|
||
codec_ = config.codec;
|
||
bitrate_ = config.bitrate;
|
||
|
||
// 对于mock URL，只初始化编码器
|
||
if (url_.substr(0, 7) == "mock://") {
|
||
if (!initFFmpeg()) {
|
||
return false;
|
||
}
|
||
is_initialized_ = true;
|
||
Logger::info("RtspWriter initialized successfully: " + name_);
|
||
return true;
|
||
}
|
||
|
||
// 对于实际URL，完整初始化
|
||
if (!initFFmpeg()) {
|
||
return false;
|
||
}
|
||
|
||
is_initialized_ = true;
|
||
Logger::info("RtspWriter initialized successfully: " + name_);
|
||
return true;
|
||
}
|
||
|
||
bool RtspWriter::write(const cv::Mat& frame) {
|
||
std::lock_guard<std::mutex> lock(mutex_);
|
||
|
||
if (!is_initialized_) {
|
||
setLastError("RtspWriter not initialized");
|
||
return false;
|
||
}
|
||
|
||
if (frame.empty()) {
|
||
setLastError("Input frame is empty");
|
||
return false;
|
||
}
|
||
|
||
try {
|
||
// 对于mock URL，只更新状态
|
||
if (url_.substr(0, 7) == "mock://") {
|
||
is_opened_ = true;
|
||
return true;
|
||
}
|
||
|
||
// 延迟打开RTSP输出，直到收到第一帧
|
||
if (!format_ctx_) {
|
||
if (!openRtspOutput()) {
|
||
return false;
|
||
}
|
||
}
|
||
|
||
// 转换图像格式
|
||
uint8_t* srcData[4] = {frame.data, nullptr, nullptr, nullptr};
|
||
int srcLinesize[4] = {static_cast<int>(frame.step), 0, 0, 0};
|
||
sws_scale(sws_ctx_, srcData, srcLinesize, 0, frame.rows,
|
||
frame_->data, frame_->linesize);
|
||
|
||
frame_->pts += av_rescale_q(1, codec_ctx_->time_base, video_stream_->time_base);
|
||
|
||
// 编码并发送
|
||
int ret = avcodec_send_frame(codec_ctx_, frame_);
|
||
if (ret < 0) {
|
||
setLastError("Error sending frame to encoder");
|
||
return false;
|
||
}
|
||
|
||
while (ret >= 0) {
|
||
ret = avcodec_receive_packet(codec_ctx_, packet_);
|
||
if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
|
||
break;
|
||
} else if (ret < 0) {
|
||
setLastError("Error receiving packet from encoder");
|
||
return false;
|
||
}
|
||
|
||
packet_->stream_index = video_stream_->index;
|
||
ret = av_interleaved_write_frame(format_ctx_, packet_);
|
||
if (ret < 0) {
|
||
setLastError("Error writing packet");
|
||
return false;
|
||
}
|
||
}
|
||
|
||
is_opened_ = true;
|
||
return true;
|
||
} catch (const std::exception& e) {
|
||
setLastError("Error writing frame: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
bool RtspWriter::isOpened() const {
|
||
std::lock_guard<std::mutex> lock(mutex_);
|
||
return is_initialized_ && is_opened_;
|
||
}
|
||
|
||
void RtspWriter::release() {
|
||
std::lock_guard<std::mutex> lock(mutex_);
|
||
cleanup();
|
||
is_initialized_ = false;
|
||
is_opened_ = false;
|
||
}
|
||
|
||
std::string RtspWriter::getName() const {
|
||
return name_;
|
||
}
|
||
|
||
bool RtspWriter::getStatus(std::string& error_msg) const {
|
||
std::lock_guard<std::mutex> lock(mutex_);
|
||
error_msg = last_error_;
|
||
return is_initialized_ && is_opened_;
|
||
}
|
||
|
||
bool RtspWriter::initFFmpeg() {
|
||
// 初始化编码器
|
||
const AVCodec* codec;
|
||
if (codec_ == "h264") {
|
||
codec = avcodec_find_encoder(AV_CODEC_ID_H264);
|
||
} else if (codec_ == "h265") {
|
||
codec = avcodec_find_encoder(AV_CODEC_ID_H265);
|
||
} else {
|
||
setLastError("Unsupported codec: " + codec_);
|
||
return false;
|
||
}
|
||
|
||
if (!codec) {
|
||
setLastError("Codec not found");
|
||
return false;
|
||
}
|
||
|
||
codec_ctx_ = avcodec_alloc_context3(codec);
|
||
if (!codec_ctx_) {
|
||
setLastError("Could not allocate codec context");
|
||
return false;
|
||
}
|
||
|
||
// 设置编码器参数
|
||
codec_ctx_->bit_rate = bitrate_;
|
||
codec_ctx_->width = 1920; // 默认分辨率，将在收到第一帧时更新
|
||
codec_ctx_->height = 1080;
|
||
codec_ctx_->time_base = (AVRational){1, fps_};
|
||
codec_ctx_->framerate = (AVRational){fps_, 1};
|
||
codec_ctx_->gop_size = 12;
|
||
codec_ctx_->pix_fmt = AV_PIX_FMT_YUV420P;
|
||
codec_ctx_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
|
||
|
||
// 对于mock URL，不需要真正打开编码器
|
||
if (url_.substr(0, 7) == "mock://") {
|
||
return true;
|
||
}
|
||
|
||
if (avcodec_open2(codec_ctx_, codec, nullptr) < 0) {
|
||
setLastError("Could not open codec");
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
return true;
|
||
}
|
||
|
||
bool RtspWriter::openRtspOutput() {
|
||
// 创建输出上下文
|
||
AVDictionary* options = nullptr;
|
||
|
||
// 设置协议白名单
|
||
av_dict_set(&options, "protocol_whitelist", "file,unix,tcp,rtsp,rtp,udp", 0);
|
||
|
||
avformat_alloc_output_context2(&format_ctx_, nullptr, "rtsp", url_.c_str());
|
||
if (!format_ctx_) {
|
||
setLastError("Could not create output context");
|
||
av_dict_free(&options);
|
||
return false;
|
||
}
|
||
|
||
// 创建视频流
|
||
video_stream_ = avformat_new_stream(format_ctx_, nullptr);
|
||
if (!video_stream_) {
|
||
setLastError("Could not create video stream");
|
||
av_dict_free(&options);
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
video_stream_->time_base = codec_ctx_->time_base;
|
||
avcodec_parameters_from_context(video_stream_->codecpar, codec_ctx_);
|
||
|
||
// 分配帧和数据包
|
||
frame_ = av_frame_alloc();
|
||
if (!frame_) {
|
||
setLastError("Could not allocate frame");
|
||
av_dict_free(&options);
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
frame_->format = codec_ctx_->pix_fmt;
|
||
frame_->width = codec_ctx_->width;
|
||
frame_->height = codec_ctx_->height;
|
||
frame_->pts = 0;
|
||
|
||
if (av_frame_get_buffer(frame_, 0) < 0) {
|
||
setLastError("Could not allocate frame data");
|
||
av_dict_free(&options);
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
packet_ = av_packet_alloc();
|
||
if (!packet_) {
|
||
setLastError("Could not allocate packet");
|
||
av_dict_free(&options);
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
// 创建图像格式转换上下文
|
||
sws_ctx_ = sws_getContext(
|
||
codec_ctx_->width, codec_ctx_->height, AV_PIX_FMT_BGR24,
|
||
codec_ctx_->width, codec_ctx_->height, codec_ctx_->pix_fmt,
|
||
SWS_BICUBIC, nullptr, nullptr, nullptr);
|
||
if (!sws_ctx_) {
|
||
setLastError("Could not create scale context");
|
||
av_dict_free(&options);
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
// 设置RTSP选项
|
||
if (url_.substr(0, 7) != "unix://") {
|
||
av_dict_set(&options, "rtsp_transport", "tcp", 0);
|
||
}
|
||
av_dict_set(&options, "muxdelay", "0.1", 0);
|
||
|
||
if (avformat_write_header(format_ctx_, &options) < 0) {
|
||
setLastError("Could not write header");
|
||
av_dict_free(&options);
|
||
cleanup();
|
||
return false;
|
||
}
|
||
|
||
av_dict_free(&options);
|
||
return true;
|
||
}
|
||
|
||
void RtspWriter::cleanup() {
|
||
if (format_ctx_) {
|
||
av_write_trailer(format_ctx_);
|
||
avio_closep(&format_ctx_->pb);
|
||
avformat_free_context(format_ctx_);
|
||
format_ctx_ = nullptr;
|
||
}
|
||
|
||
if (codec_ctx_) {
|
||
avcodec_free_context(&codec_ctx_);
|
||
}
|
||
|
||
if (frame_) {
|
||
av_frame_free(&frame_);
|
||
}
|
||
|
||
if (packet_) {
|
||
av_packet_free(&packet_);
|
||
}
|
||
|
||
if (sws_ctx_) {
|
||
sws_freeContext(sws_ctx_);
|
||
sws_ctx_ = nullptr;
|
||
}
|
||
}
|
||
|
||
void RtspWriter::setLastError(const std::string& error) {
|
||
last_error_ = error;
|
||
Logger::error("RtspWriter error: " + error);
|
||
}
|
||
|
||
} // namespace pipeline
|