140 lines
5.3 KiB
Python
140 lines
5.3 KiB
Python
import os
|
||
import cv2
|
||
import re
|
||
from surround_view import CaptureThread, CameraProcessingThread
|
||
from surround_view import FisheyeCameraModel, BirdView
|
||
from surround_view import MultiBufferManager, ProjectedImageBuffer
|
||
import surround_view.param_settings as settings
|
||
|
||
|
||
def is_rtsp_url(input_str):
|
||
"""检查输入是否为RTSP URL"""
|
||
return input_str.startswith('rtsp://')
|
||
|
||
|
||
def main():
|
||
# 配置四个RTSP摄像头地址（前、后、左、右）
|
||
camera_ids = [
|
||
"rtsp://admin:@192.168.112.153:554/video", # 前
|
||
"rtsp://admin:@192.168.112.152:554/video", # 后
|
||
"rtsp://admin:@192.168.112.150:554/video", # 左
|
||
"rtsp://admin:@192.168.112.151:554/video" # 右
|
||
]
|
||
|
||
# 摄像头ID映射（使用索引作为设备ID）
|
||
camera_id_mapping = {i: cam_id for i, cam_id in enumerate(camera_ids)}
|
||
|
||
# 摄像头翻转参数（根据实际安装方向调整）
|
||
flip_methods = [0, 0, 0, 0] # 0:不翻转, 1:水平翻转, 2:垂直翻转, 3:水平+垂直
|
||
|
||
# 加载相机内参模型
|
||
yamls_dir = os.path.join(os.getcwd(), "yaml")
|
||
camera_files = [os.path.join(yamls_dir, name + ".yaml") for name in settings.camera_names]
|
||
camera_models = [FisheyeCameraModel(camera_file, name)
|
||
for camera_file, name in zip(camera_files, settings.camera_names)]
|
||
|
||
# 初始化捕获线程（针对RTSP特殊处理）
|
||
capture_tds = []
|
||
for idx, (cam_id, flip) in enumerate(zip(camera_ids, flip_methods)):
|
||
# 使用索引作为设备ID，解决KeyError问题
|
||
if is_rtsp_url(cam_id):
|
||
# RTSP流配置（使用FFmpeg后端）
|
||
capture_tds.append(CaptureThread(
|
||
device_id=idx, # 使用索引作为设备ID
|
||
flip_method=flip,
|
||
use_gst=False, # 关闭GStreamer，使用FFmpeg
|
||
api_preference=cv2.CAP_FFMPEG, # 强制使用FFmpeg
|
||
resolution=(960, 640)
|
||
))
|
||
# 手动打开RTSP流
|
||
if not capture_tds[-1].cap.open(cam_id, cv2.CAP_FFMPEG):
|
||
print(f"无法打开RTSP流: {cam_id}")
|
||
return
|
||
# 设置RTSP参数（兼容不同OpenCV版本）
|
||
try:
|
||
# 尝试通用参数设置方式
|
||
capture_tds[-1].cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) # 减少缓冲
|
||
capture_tds[-1].cap.set(cv2.CAP_PROP_FRAME_WIDTH, 960)
|
||
capture_tds[-1].cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 640)
|
||
# 处理FFMPEG选项（兼容旧版本）
|
||
if hasattr(cv2, 'CAP_PROP_FFMPEG_OPTION'):
|
||
capture_tds[-1].cap.set(cv2.CAP_PROP_FFMPEG_OPTION, "rtsp_transport", "tcp")
|
||
capture_tds[-1].cap.set(cv2.CAP_PROP_FFMPEG_OPTION, "stimeout", "5000000")
|
||
else:
|
||
print(f"RTSP流 {cam_id} 使用兼容模式，部分参数可能无法设置")
|
||
except Exception as e:
|
||
print(f"设置RTSP参数警告: {e}")
|
||
else:
|
||
# 普通USB/CSI相机配置
|
||
capture_tds.append(CaptureThread(
|
||
device_id=idx, # 使用索引作为设备ID
|
||
flip_method=flip,
|
||
use_gst=True,
|
||
resolution=(960, 640)
|
||
))
|
||
if not capture_tds[-1].connect_camera():
|
||
print(f"启动摄像头 {cam_id} 失败，退出程序")
|
||
return
|
||
|
||
# 绑定捕获缓冲区管理器
|
||
capture_buffer_manager = MultiBufferManager()
|
||
for td in capture_tds:
|
||
capture_buffer_manager.bind_thread(td, buffer_size=8)
|
||
td.start()
|
||
|
||
# 初始化处理线程（使用索引作为device_id）
|
||
proc_buffer_manager = ProjectedImageBuffer()
|
||
process_tds = [
|
||
CameraProcessingThread(
|
||
capture_buffer_manager,
|
||
device_id=idx, # 使用索引作为设备ID
|
||
camera_model=model
|
||
)
|
||
for idx, model in enumerate(camera_models)
|
||
]
|
||
|
||
# 启动处理线程
|
||
for td in process_tds:
|
||
proc_buffer_manager.bind_thread(td)
|
||
td.start()
|
||
|
||
# 初始化环视拼接
|
||
birdview = BirdView(proc_buffer_manager)
|
||
birdview.load_weights_and_masks("./weights.png", "./masks.png")
|
||
birdview.start()
|
||
|
||
# 主循环显示
|
||
try:
|
||
while True:
|
||
birdview_img = birdview.get()
|
||
if birdview_img is not None:
|
||
display_img = cv2.resize(birdview_img, (800, 600))
|
||
cv2.imshow("Surround View", display_img)
|
||
|
||
key = cv2.waitKey(1) & 0xFF
|
||
if key == ord("q"):
|
||
break
|
||
|
||
# 打印帧率信息
|
||
status = []
|
||
for td in capture_tds:
|
||
cam_id = camera_id_mapping[td.device_id]
|
||
status.append(f"捕获 {cam_id.split('@')[-1].split(':')[0]} FPS: {td.stat_data.average_fps:.1f}")
|
||
for td in process_tds:
|
||
cam_name = settings.camera_names[td.device_id]
|
||
status.append(f"处理 {cam_name} FPS: {td.stat_data.average_fps:.1f}")
|
||
status.append(f"全景 FPS: {birdview.stat_data.average_fps:.1f}")
|
||
print(" | ".join(status), end="\r")
|
||
|
||
finally:
|
||
for td in process_tds:
|
||
td.stop()
|
||
for td in capture_tds:
|
||
td.stop()
|
||
td.disconnect_camera()
|
||
birdview.stop()
|
||
cv2.destroyAllWindows()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |