主要改动
1. 函数签名更新
添加了推流相关参数:
13:22:infer.py
def process_video_stream(video_url, conf_threshold=0.5, show_display=True, save_output=False, save_coco=True,
stream_output=False, stream_url=None):
"""
使用YOLOv8对视频流进行目标检测
参数:
video_url: 视频流地址(输入)
conf_threshold: 置信度阈值
show_display: 是否显示实时画面
save_output: 是否保存处理后的视频到本地文件
save_coco: 是否保存COCO格式的检测结果
stream_output: 是否推流到服务器(RTMP/RTSP等)
stream_url: 推流地址(如:rtmp://server.com/live/stream_key)
"""
2. 推流配置验证
107:115:infer.py
# 验证推流配置
if stream_output:
if not stream_url:
print("❌ 推流已启用但未提供推流地址,将禁用推流")
stream_output = False
elif not stream_url.startswith(('rtmp://', 'rtsp://', 'http://', 'https://')):
print(f"⚠️ 警告: 推流地址格式可能不正确: {stream_url}")
print(" 支持的格式: rtmp://, rtsp://, http://, https://")
3. 推流进程初始化
使用 ffmpeg 进行低延迟推流:
386:420:infer.py
# 设置推流输出
if stream_output and stream_url:
print(f"\n📡 设置实时推流...")
# 判断推流协议类型
if stream_url.startswith('rtmp://'):
stream_format = 'flv'
elif stream_url.startswith('rtsp://'):
stream_format = 'rtsp'
else:
stream_format = 'flv' # 默认使用flv
# 构建ffmpeg推流命令(低延迟配置)
stream_cmd = [
'ffmpeg',
'-f', 'rawvideo',
'-vcodec', 'rawvideo',
'-s', f'{width}x{height}',
'-pix_fmt', 'bgr24',
'-r', str(fps),
'-i', '-', # 从stdin读取
'-an', # 无音频
'-vcodec', 'libx264',
'-preset', 'ultrafast', # 超快编码,低延迟
'-tune', 'zerolatency', # 零延迟调优
'-pix_fmt', 'yuv420p',
'-f', stream_format,
'-flvflags', 'no_duration_filesize', # RTMP优化
stream_url
]
4. 推流写入逻辑
在每一帧处理后写入推流:
765:800:infer.py
# 推流到服务器(如果需要)
if stream_output and stream_process is not None:
try:
# 使用相同的annotated_frame(已经处理过格式)
# 确保帧尺寸匹配
frame_h, frame_w = annotated_frame.shape[:2]
if frame_w != width or frame_h != height:
stream_frame = cv2.resize(annotated_frame, (width, height), interpolation=cv2.INTER_LINEAR)
else:
stream_frame = annotated_frame.copy()
# 确保帧格式正确
if len(stream_frame.shape) != 3 or stream_frame.shape[2] != 3:
if frame_count % 100 == 0:
print(f"⚠️ 警告: 第 {frame_count} 帧格式异常,跳过推流")
else:
# 确保数据类型正确(uint8)
if stream_frame.dtype != np.uint8:
if stream_frame.max() > 255 or stream_frame.min() < 0:
stream_frame = np.clip(stream_frame, 0, 255)
stream_frame = stream_frame.astype(np.uint8)
# 确保帧是连续的(C-contiguous)
if not stream_frame.flags['C_CONTIGUOUS']:
stream_frame = np.ascontiguousarray(stream_frame)
# 写入到推流进程的stdin
try:
stream_process.stdin.write(stream_frame.tobytes())
if frame_count % 100 == 0:
print(f"📡 已推流 {frame_count} 帧")
except (BrokenPipeError, OSError) as e:
if frame_count % 10 == 0:
print(f"⚠️ 警告: 推流写入失败: {e}")
print(f" 推流可能已断开,请检查推流服务器状态")
except Exception as e:
if frame_count % 10 == 0:
print(f"⚠️ 警告: 推流处理异常: {e}")
5. 推流进程清理
在 finally 块中正确关闭推流:
808:820:infer.py
# 关闭推流进程
if stream_output and stream_process is not None:
try:
print("\n📡 正在关闭推流...")
stream_process.stdin.close()
stream_process.wait(timeout=5) # 等待推流完成
print("✓ 推流已关闭")
except subprocess.TimeoutExpired:
print("⚠️ 推流关闭超时,强制终止")
stream_process.kill()
except Exception as e:
print(f"⚠️ 关闭推流时出错: {e}")
6. 配置示例
在 main 函数中添加了配置示例:
850:862:infer.py
# 配置参数
config = {
'conf_threshold': 0.5, # 置信度阈值 (0-1之间,越高检测越严格)
'show_display': True, # 是否显示实时画面
'save_output': True, # 是否保存处理后的视频到本地文件
'save_coco': True, # 是否保存COCO格式的检测结果
'stream_output': False, # 是否推流到服务器(RTMP/RTSP等)
'stream_url': None # 推流地址,例如: rtmp://server.com/live/stream_key
}
# 推流配置示例(取消注释以启用)
# config['stream_output'] = True
# config['stream_url'] = 'rtmp://your-server.com/live/stream_key'
功能特点
- 开关配置:通过
stream_output控制是否推流 - 低延迟:使用
ultrafast和zerolatency参数 - 多协议支持:支持 RTMP、RTSP 等
- 独立进程:推流和保存文件可同时进行
- 错误处理:包含推流断开的检测和处理
使用方法
启用推流只需在配置中设置:
python
config = {
'stream_output': True, # 启用推流
'stream_url': 'rtmp://your-server.com/live/stream_key' # 推流地址
}
代码已支持实时推流功能,可通过配置开关控制。