apps/generic-rtsp-yolov8-render
1
0
Fork 0
mirror of https://github.com/wwhai/generic-rtsp-yolov8-render.git synced 2025-12-24 12:12:50 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

dev: update

Browse Source
This commit is contained in:
wwhai
2025-02-07 23:46:42 +08:00
parent b0bce76c2b
commit 0ab215433c
8 changed files with 137 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
include/libav_utils.h
View File

@@ -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);

7
include/push_stream_thread.h
View File

@@ -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

5
include/thread_args.h
View File

@@ -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;

2
readme.md
View File

@@ -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 环境

31
src/libav_utils.cc
View File

@@ -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)
{

43
src/main.cc
View File

@@ -25,13 +25,11 @@
#include <unistd.h>
#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 <RTSP_URL>\n", argv[0]);
fprintf(stderr, "Usage: %s <push_to_rtsp_url>\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);

191
src/push_stream_thread.cc
View File

@@ -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;

34
src/rtsp_handler.cc
View File

@@ -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;