From 0ab215433cb8dada72007b7d2e553bd211f22d19 Mon Sep 17 00:00:00 2001 From: wwhai Date: Fri, 7 Feb 2025 23:46:42 +0800 Subject: [PATCH] dev: update --- include/libav_utils.h | 5 - include/push_stream_thread.h | 7 +- include/thread_args.h | 5 +- readme.md | 2 +- src/libav_utils.cc | 31 ------ src/main.cc | 43 +++----- src/push_stream_thread.cc | 191 ++++++++++++++++------------------- src/rtsp_handler.cc | 34 +++++-- 8 files changed, 137 insertions(+), 181 deletions(-) diff --git a/include/libav_utils.h b/include/libav_utils.h index 9182abf..59548f5 100644 --- a/include/libav_utils.h +++ b/include/libav_utils.h @@ -27,11 +27,6 @@ extern "C" #include "frame_queue.h" // 获取错误字符串的全局缓冲区实现 const char *get_av_error(int errnum); -// 函数用于复制AVFrame -// 注意：使用完毕后需要释放 -// @param srcFrame 源帧 -// @return 复制后的帧 -AVFrame *CopyAVFrame(AVFrame *srcFrame); // 截图 int CaptureImage(AVFrame *srcFrame, const char *file_path); diff --git a/include/push_stream_thread.h b/include/push_stream_thread.h index e94fead..5545b4d 100644 --- a/include/push_stream_thread.h +++ b/include/push_stream_thread.h @@ -35,9 +35,10 @@ extern "C" #include "thread_args.h" typedef struct { - AVFormatContext *output_ctx; - AVStream *video_stream; - AVCodecContext *codec_ctx; + AVFormatContext *output_ctx; // 输出格式上下文 + AVCodecContext *output_codec_ctx; // 输出编码器上下文 + AVStream *input_stream; // 输入流 + AVStream *output_stream; // 输出流 } RtmpStreamContext; /// @brief diff --git a/include/thread_args.h b/include/thread_args.h index 3e90c18..8bd237d 100644 --- a/include/thread_args.h +++ b/include/thread_args.h @@ -20,12 +20,15 @@ typedef struct { - const char *rtsp_url; + const char *input_stream_url; + const char *output_stream_url; FrameQueue *video_queue; FrameQueue *detection_queue; FrameQueue *box_queue; FrameQueue *origin_frame_queue; FrameQueue *infer_frame_queue; + AVStream *input_stream; // 输入流 + AVStream *output_stream; // 输出流 Context *ctx; } ThreadArgs; diff --git a/readme.md b/readme.md index 1a23996..92e8afa 100644 --- a/readme.md +++ b/readme.md @@ -32,7 +32,7 @@ make clean 编译成功后，运行以下命令启动程序： ```bash -./generic-rtsp-yolov8-render.exe rtsp://192.168.10.8:554/av0_0 +./generic-rtsp-yolov8-render.exe rtsp://192.168.10.8:554/av0_0 rtmp://192.168.10.9:1935/live/tlive001 ``` ### Docker 环境 diff --git a/src/libav_utils.cc b/src/libav_utils.cc index 1b08591..7bd4c31 100644 --- a/src/libav_utils.cc +++ b/src/libav_utils.cc @@ -21,38 +21,7 @@ const char *get_av_error(int errnum) av_strerror(errnum, str, sizeof(str)); return str; } -// 函数用于复制AVFrame -AVFrame *CopyAVFrame(AVFrame *srcFrame) -{ - AVFrame *dstFrame = av_frame_alloc(); - if (!dstFrame) - { - fprintf(stderr, "Could not allocate frame\n"); - return NULL; - } - int ret = av_frame_get_buffer(dstFrame, 32); - if (ret < 0) - { - fprintf(stderr, "Could not allocate new frame buffer\n"); - av_frame_free(&dstFrame); - return NULL; - } - ret = av_frame_copy(dstFrame, srcFrame); - if (ret < 0) - { - fprintf(stderr, "Could not copy frame\n"); - av_frame_free(&dstFrame); - return NULL; - } - if (av_frame_copy_props(dstFrame, srcFrame) < 0) - { - fprintf(stderr, "Could not copy frame properties\n"); - av_frame_free(&dstFrame); - return NULL; - } - return dstFrame; -} // 保存图像到文件,格式为png，不要用Libav的库，直接提取Avframe的Data int CaptureImage(AVFrame *srcFrame, const char *file_path) { diff --git a/src/main.cc b/src/main.cc index 0f692f8..7d9fee4 100644 --- a/src/main.cc +++ b/src/main.cc @@ -25,13 +25,11 @@ #include #include "background.h" #include "context.h" -#include "push_stream_thread.h" /// @brief Context *background_thread_ctx; Context *pull_rtsp_thread_ctx; Context *video_renderer_thread_ctx; Context *detection_thread_ctx; -Context *push_rtsp_thread_ctx; /// @brief /// @param sig void handle_signal(int sig) @@ -40,16 +38,15 @@ void handle_signal(int sig) CancelContext(video_renderer_thread_ctx); CancelContext(detection_thread_ctx); CancelContext(background_thread_ctx); - CancelContext(push_rtsp_thread_ctx); fprintf(stderr, "Received signal %d, exiting...\n", sig); exit(0); } int main(int argc, char *argv[]) { - if (argc < 2) + if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); + fprintf(stderr, "Usage: %s \n", argv[0]); return 1; } if (signal(SIGINT, handle_signal) == SIG_ERR) @@ -57,13 +54,15 @@ int main(int argc, char *argv[]) perror("Failed to set signal handler for SIGTERM"); return 1; } - const char *rtsp_url = argv[1]; + const char *pull_from_rtsp_url = argv[1]; + const char *push_to_rtsp_url = argv[2]; + fprintf(stderr, "pull_from_rtsp_url: %s\n", pull_from_rtsp_url); + fprintf(stderr, "push_to_rtsp_url: %s\n", push_to_rtsp_url); + // background_thread_ctx = CreateContext(); pull_rtsp_thread_ctx = CreateContext(); video_renderer_thread_ctx = CreateContext(); detection_thread_ctx = CreateContext(); - push_rtsp_thread_ctx = CreateContext(); - // Initialize frame queues FrameQueue video_queue, detection_queue, box_queue, origin_frame_queue, infer_frame_queue; frame_queue_init(&video_queue, 60); @@ -73,22 +72,20 @@ int main(int argc, char *argv[]) frame_queue_init(&infer_frame_queue, 60); // Create threads - pthread_t background_thread, pull_rtsp_thread, renderer_thread, detection_thread, push_rtsp_thread; + pthread_t background_thread, pull_rtsp_thread, renderer_thread, detection_thread; // ThreadArgs background_thread_args = {.ctx = background_thread_ctx}; - ThreadArgs pull_rtsp_thread_args = {rtsp_url, &video_queue, + ThreadArgs pull_rtsp_thread_args = {pull_from_rtsp_url, push_to_rtsp_url, &video_queue, &detection_queue, &box_queue, &origin_frame_queue, - &infer_frame_queue, pull_rtsp_thread_ctx}; - ThreadArgs video_renderer_thread_args = {rtsp_url, &video_queue, + &infer_frame_queue, NULL, NULL, pull_rtsp_thread_ctx}; + ThreadArgs video_renderer_thread_args = {pull_from_rtsp_url, push_to_rtsp_url, &video_queue, &detection_queue, &box_queue, &origin_frame_queue, - &infer_frame_queue, pull_rtsp_thread_ctx}; - ThreadArgs detection_thread_args = {rtsp_url, &video_queue, + &infer_frame_queue, NULL, NULL, video_renderer_thread_ctx}; + ThreadArgs detection_thread_args = {pull_from_rtsp_url, push_to_rtsp_url, &video_queue, &detection_queue, &box_queue, &origin_frame_queue, - &infer_frame_queue, pull_rtsp_thread_ctx}; - ThreadArgs push_rtsp_thread_args = {rtsp_url, &video_queue, - &detection_queue, &box_queue, &origin_frame_queue, - &infer_frame_queue, pull_rtsp_thread_ctx}; + &infer_frame_queue, NULL, NULL, detection_thread_ctx}; + // if (pthread_create(&background_thread, NULL, background_task_thread, (void *)&background_thread_args) != 0) { @@ -113,17 +110,11 @@ int main(int argc, char *argv[]) perror("Failed to create detection thread"); goto END; } - // - if (pthread_create(&push_rtsp_thread, NULL, push_rtmp_handler_thread, (void *)&push_rtsp_thread_args) != 0) - { - perror("Failed to create detection thread"); - goto END; - } + fprintf(stderr, "Main thread waiting for threads to finish...\n"); pthread_detach(pull_rtsp_thread); pthread_detach(renderer_thread); pthread_detach(detection_thread); - pthread_detach(push_rtsp_thread); pthread_join(background_thread, NULL); END: // Cancel contexts @@ -131,13 +122,11 @@ END: CancelContext(pull_rtsp_thread_ctx); CancelContext(video_renderer_thread_ctx); CancelContext(detection_thread_ctx); - CancelContext(push_rtsp_thread_ctx); // Destroy threads pthread_cond_destroy(&background_thread_ctx->cond); pthread_mutex_destroy(&pull_rtsp_thread_ctx->mtx); pthread_mutex_destroy(&video_renderer_thread_ctx->mtx); pthread_mutex_destroy(&detection_thread_ctx->mtx); - pthread_mutex_destroy(&push_rtsp_thread_ctx->mtx); // Free frame queues frame_queue_destroy(&video_queue); frame_queue_destroy(&box_queue); diff --git a/src/push_stream_thread.cc b/src/push_stream_thread.cc index 35bc453..7223f15 100644 --- a/src/push_stream_thread.cc +++ b/src/push_stream_thread.cc @@ -18,72 +18,64 @@ // 初始化 RTMP 流上下文 int init_rtmp_stream(RtmpStreamContext *ctx, const char *output_url, int width, int height, int fps) { - // 输出参数 - fprintf(stderr, "init_rtmp_stream === output_url=%s,width=%d,height=%d,fps=%d\n", - output_url, width, height, fps); - // 创建输出上下文 - int ret = avformat_alloc_output_context2(&ctx->output_ctx, NULL, "flv", output_url); + int ret; + + // 1. 打开输出流上下文 + ret = avformat_alloc_output_context2(&ctx->output_ctx, NULL, "flv", output_url); if (ret < 0 || !ctx->output_ctx) { - fprintf(stderr, "Failed to create output context: %s\n", get_av_error(ret)); + fprintf(stderr, "Failed to allocate output context: %s\n", get_av_error(ret)); return -1; } - // 查找编码器 + // 2. 配置视频编码器 const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264); if (!codec) { - fprintf(stderr, "H.264 encoder not found\n"); + fprintf(stderr, "H264 encoder not found\n"); return -1; } - // 创建编码器上下文 - ctx->codec_ctx = avcodec_alloc_context3(codec); - if (!ctx->codec_ctx) - { - fprintf(stderr, "Failed to allocate codec context\n"); - return -1; - } - - // 配置编码参数 - ctx->codec_ctx->width = width; - ctx->codec_ctx->height = height; - ctx->codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P; - ctx->codec_ctx->time_base = (AVRational){1, fps}; - ctx->codec_ctx->framerate = (AVRational){fps, 1}; - ctx->codec_ctx->bit_rate = 4000000; - ctx->codec_ctx->gop_size = 12; - - // H.264高级配置 - // av_opt_set(ctx->codec_ctx->priv_data, "preset", "fast", 0); - // av_opt_set(ctx->codec_ctx->priv_data, "tune", "zerolatency", 0); - - // 创建输出流 - ctx->video_stream = avformat_new_stream(ctx->output_ctx, NULL); - if (!ctx->video_stream) + ctx->output_stream = avformat_new_stream(ctx->output_ctx, codec); + if (!ctx->output_stream) { fprintf(stderr, "Failed to create video stream\n"); return -1; } - // 关联编码器参数到流 - ret = avcodec_parameters_from_context(ctx->video_stream->codecpar, ctx->codec_ctx); - if (ret < 0) + ctx->output_codec_ctx = avcodec_alloc_context3(codec); + if (!ctx->output_codec_ctx) + { + fprintf(stderr, "Failed to allocate codec context\n"); + return -1; + } + if (avcodec_parameters_copy(ctx->output_stream->codecpar, ctx->input_stream->codecpar) < 0) { fprintf(stderr, "Failed to copy codec parameters: %s\n", get_av_error(ret)); return -1; } - - // 打开编码器 - ret = avcodec_open2(ctx->codec_ctx, codec, NULL); + if (ctx->output_ctx->oformat->flags & AVFMT_GLOBALHEADER) + { + ctx->output_codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER; + } + // 4. 打开编码器 + ret = avcodec_open2(ctx->output_codec_ctx, codec, NULL); if (ret < 0) { fprintf(stderr, "Failed to open codec: %s\n", get_av_error(ret)); return -1; } - // 打开网络输出 - if (!(ctx->output_ctx->oformat->flags & AVFMT_NOFILE)) + // 5. 设置流编码器参数 + ret = avcodec_parameters_from_context(ctx->output_stream->codecpar, ctx->output_codec_ctx); + if (ret < 0) + { + fprintf(stderr, "Failed to copy codec parameters: %s\n", get_av_error(ret)); + return -1; + } + + // 6. 打开输出流 + if (!(ctx->output_ctx->flags & AVFMT_NOFILE)) { ret = avio_open(&ctx->output_ctx->pb, output_url, AVIO_FLAG_WRITE); if (ret < 0) @@ -93,107 +85,98 @@ int init_rtmp_stream(RtmpStreamContext *ctx, const char *output_url, int width, } } - // 写入文件头 + // 7. 写文件头 ret = avformat_write_header(ctx->output_ctx, NULL); if (ret < 0) { - fprintf(stderr, "Failed to write header:%d, %s\n", ret, get_av_error(ret)); + fprintf(stderr, "Failed to write header: %s\n", get_av_error(ret)); return -1; } return 0; } -// 优化后的push_stream函数 +// 在推送帧之前，确保 PTS 和 DTS 是递增的 +static int64_t last_pts = 0; +static int64_t last_dts = 0; + void push_stream(RtmpStreamContext *ctx, AVFrame *frame) { - if (!ctx || !frame) + int ret; + + // 设置 PTS 和 DTS + if (frame->pts == AV_NOPTS_VALUE) { - fprintf(stderr, "Invalid input parameters: RtmpStreamContext or AVFrame is NULL\n"); - return; + // 如果 PTS 无效，根据上一帧的 PTS 递增 + frame->pts = last_pts + 1; } - int ret = 0; - // 发送帧到编码器 - ret = avcodec_send_frame(ctx->codec_ctx, frame); + if (frame->pkt_dts == AV_NOPTS_VALUE) + { + // 如果 DTS 无效，根据上一帧的 DTS 递增 + frame->pkt_dts = last_dts + 1; + } + + // 更新 last_pts 和 last_dts + last_pts = frame->pts; + last_dts = frame->pkt_dts; + + // 1. 编码帧 + ret = avcodec_send_frame(ctx->output_codec_ctx, frame); if (ret < 0) { - fprintf(stderr, "Error sending frame: %s\n", get_av_error(ret)); + fprintf(stderr, "Failed to send frame for encoding: %s\n", get_av_error(ret)); return; } + // 2. 创建一个新的 AVPacket AVPacket *pkt = av_packet_alloc(); if (!pkt) { - fprintf(stderr, "Error allocating AVPacket\n"); + fprintf(stderr, "Failed to allocate AVPacket\n"); return; } - // 接收编码后的数据包 - while (1) + // 3. 获取编码后的数据包 + ret = avcodec_receive_packet(ctx->output_codec_ctx, pkt); + if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) { - ret = avcodec_receive_packet(ctx->codec_ctx, pkt); - if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) - { - // 没有更多数据或需要更多输入帧 - break; - } - else if (ret < 0) - { - fprintf(stderr, "Error encoding frame: %s\n", get_av_error(ret)); - break; - } - // 输出 frame index - fprintf(stderr, "Packet index: %d\n", pkt->stream_index); - fprintf(stderr, "Codec time base: %d/%d\n", ctx->codec_ctx->time_base.num, ctx->codec_ctx->time_base.den); - fprintf(stderr, "Video stream time base: %d/%d\n", ctx->video_stream->time_base.num, ctx->video_stream->time_base.den); - fprintf(stderr, "BEFORE ====== PTS: %ld, DTS: %ld, Duration: %ld\n", pkt->pts, pkt->dts, pkt->duration); - - // 调整时间戳 - pkt->pts = av_rescale_q_rnd(pkt->pts, ctx->codec_ctx->time_base, ctx->video_stream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX)); - pkt->dts = av_rescale_q_rnd(pkt->dts, ctx->codec_ctx->time_base, ctx->video_stream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX)); - int fps = 25; - AVRational codec_time_base = ctx->codec_ctx->time_base; - pkt->duration = av_rescale_q(1, (AVRational){1, fps}, codec_time_base); - pkt->pos = -1; - pkt->stream_index = ctx->video_stream->index; - - fprintf(stderr, "AFTER ====== PTS: %ld, DTS: %ld, Duration: %ld\n", pkt->pts, pkt->dts, pkt->duration); - - // 检查数据包大小 - if (pkt->size <= 0) - { - fprintf(stderr, "Packet size is invalid: %d\n", pkt->size); - return; - } - ret = av_interleaved_write_frame(ctx->output_ctx, pkt); - if (ret < 0) - { - fprintf(stderr, "Error writing packet: %d, %s\n", ret, get_av_error(ret)); - } - - // 释放数据包 - av_packet_unref(pkt); + av_packet_free(&pkt); + return; + } + else if (ret < 0) + { + fprintf(stderr, "Failed to receive packet: %s\n", get_av_error(ret)); + av_packet_free(&pkt); + return; } + // 4. 设置时间戳 + pkt->stream_index = ctx->output_stream->index; + av_packet_rescale_ts(pkt, ctx->output_codec_ctx->time_base, ctx->output_stream->time_base); + + // 5. 推送数据包 + ret = av_write_frame(ctx->output_ctx, pkt); + if (ret < 0) + { + fprintf(stderr, "Failed to write frame: %s\n", get_av_error(ret)); + } + + // 6. 释放 AVPacket av_packet_free(&pkt); } - // 推流线程处理函数 void *push_rtmp_handler_thread(void *arg) { ThreadArgs *args = (ThreadArgs *)arg; - const char *output_url = "rtmp://192.168.10.8:1935/live/tlive001"; - RtmpStreamContext ctx; memset(&ctx, 0, sizeof(RtmpStreamContext)); - - // 初始化输出流 - if (init_rtmp_stream(&ctx, output_url, 1920, 1080, 25) < 0) + fprintf(stderr, "push_rtmp_handler_thread\n"); + ctx.input_stream = args->input_stream; + if (init_rtmp_stream(&ctx, args->output_stream_url, 1920, 1080, 25) < 0) { fprintf(stderr, "Failed to initialize RTMP stream\n"); return NULL; } - // Main processing loop while (1) { @@ -210,9 +193,9 @@ void *push_rtmp_handler_thread(void *arg) } } } - + CancelContext(args->ctx); av_write_trailer(ctx.output_ctx); - avcodec_free_context(&ctx.codec_ctx); + avcodec_free_context(&ctx.output_codec_ctx); avio_closep(&ctx.output_ctx->pb); avformat_free_context(ctx.output_ctx); return NULL; diff --git a/src/rtsp_handler.cc b/src/rtsp_handler.cc index 3676454..5fca2df 100644 --- a/src/rtsp_handler.cc +++ b/src/rtsp_handler.cc @@ -26,7 +26,7 @@ extern "C" #include "frame_queue.h" #include "rtsp_handler.h" #include "libav_utils.h" - +#include "push_stream_thread.h" void *pull_rtsp_handler_thread(void *arg) { const ThreadArgs *args = (ThreadArgs *)arg; @@ -34,9 +34,10 @@ void *pull_rtsp_handler_thread(void *arg) int ret; // Open RTSP input stream - if ((ret = avformat_open_input(&fmt_ctx, args->rtsp_url, NULL, NULL)) < 0) + fprintf(stderr, "pull_rtsp_handler_thread: %s\n", args->input_stream_url); + if ((ret = avformat_open_input(&fmt_ctx, args->input_stream_url, NULL, NULL)) < 0) { - fprintf(stderr, "Error: Could not open RTSP stream :(%s).\n", args->rtsp_url); + fprintf(stderr, "Error: Could not open RTSP stream :(%s).\n", args->input_stream_url); pthread_exit(NULL); } @@ -66,7 +67,7 @@ void *pull_rtsp_handler_thread(void *arg) } // Print stream information - av_dump_format(fmt_ctx, 0, args->rtsp_url, 0); + av_dump_format(fmt_ctx, 0, args->input_stream_url, 0); // Allocate AVPacket for reading frames AVPacket *origin_packet = av_packet_alloc(); @@ -133,8 +134,22 @@ void *pull_rtsp_handler_thread(void *arg) printf(" codec_type: %s\n", av_get_media_type_string(stream->codecpar->codec_type)); printf(" codec_name: %s\n", avcodec_get_name(stream->codecpar->codec_id)); } - fprintf(stderr, "RTSP handler thread started. Pull stream: %s\n", args->rtsp_url); - + // Create a thread for pushing frames to the output stream + pthread_t push_rtsp_thread; + Context *push_rtmp_handler_thread_ctx = CreateContext(); + ThreadArgs push_rtmp_handler_thread_args; + // Copy the necessary data to the push_rtmp_handler_thread_args struct + push_rtmp_handler_thread_args.ctx = push_rtmp_handler_thread_ctx; + push_rtmp_handler_thread_args.origin_frame_queue = args->origin_frame_queue; + push_rtmp_handler_thread_args.output_stream_url = args->output_stream_url; + push_rtmp_handler_thread_args.input_stream = fmt_ctx->streams[video_stream_index]; + fprintf(stderr, "push_rtmp_handler_thread: %s\n", args->output_stream_url); + if (pthread_create(&push_rtsp_thread, NULL, push_rtmp_handler_thread, (void *)&push_rtmp_handler_thread_args) != 0) + { + perror("Failed to create push_rtmp_handler_thread thread"); + exit(1); + } + pthread_detach(push_rtsp_thread); // Read frames from the stream while (av_read_frame(fmt_ctx, origin_packet) >= 0) { @@ -186,7 +201,7 @@ void *pull_rtsp_handler_thread(void *arg) else { { - AVFrame *display_frame = CopyAVFrame(origin_frame); + AVFrame *display_frame = av_frame_clone(origin_frame); QueueItem outputItem; outputItem.type = ONLY_FRAME; outputItem.data = display_frame; @@ -197,7 +212,7 @@ void *pull_rtsp_handler_thread(void *arg) } } { - AVFrame *output_frame = CopyAVFrame(origin_frame); + AVFrame *output_frame = av_frame_clone(origin_frame); QueueItem outputItem; outputItem.type = ONLY_FRAME; outputItem.data = output_frame; @@ -208,7 +223,8 @@ void *pull_rtsp_handler_thread(void *arg) } } { - AVFrame *detection_frame = CopyAVFrame(origin_frame); + + AVFrame *detection_frame = av_frame_clone(origin_frame); QueueItem outputItem; outputItem.type = ONLY_FRAME; outputItem.data = detection_frame;