118 lines
3.2 KiB
Python
118 lines
3.2 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
简单的 TensorRT 推理测试
|
|||
|
|
验证 TensorRT 引擎是否能正常工作
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import yaml
|
|||
|
|
import time
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_tensorrt_inference():
|
|||
|
|
"""测试 TensorRT 推理"""
|
|||
|
|
print("TensorRT 推理测试")
|
|||
|
|
print("=" * 60)
|
|||
|
|
|
|||
|
|
# 配置
|
|||
|
|
config_path = "config.yaml"
|
|||
|
|
engine_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.engine"
|
|||
|
|
|
|||
|
|
# 加载配置
|
|||
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|||
|
|
cfg = yaml.safe_load(f)
|
|||
|
|
|
|||
|
|
# 加载 TensorRT 引擎
|
|||
|
|
print(f"🚀 加载 TensorRT 引擎: {engine_path}")
|
|||
|
|
model = YOLO(engine_path, task='detect')
|
|||
|
|
print("✅ 引擎加载成功")
|
|||
|
|
|
|||
|
|
# 获取第一个摄像头
|
|||
|
|
cam_cfg = cfg['cameras'][0]
|
|||
|
|
cam_id = cam_cfg['id']
|
|||
|
|
rtsp_url = cam_cfg['rtsp_url']
|
|||
|
|
|
|||
|
|
print(f"\n📹 连接摄像头: {cam_id}")
|
|||
|
|
print(f" RTSP: {rtsp_url}")
|
|||
|
|
|
|||
|
|
# 打开视频流
|
|||
|
|
cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG)
|
|||
|
|
|
|||
|
|
if not cap.isOpened():
|
|||
|
|
print("❌ 无法打开视频流")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
print("✅ 视频流已连接")
|
|||
|
|
|
|||
|
|
# 测试推理
|
|||
|
|
print(f"\n🔄 开始推理测试（10秒）...")
|
|||
|
|
|
|||
|
|
frame_count = 0
|
|||
|
|
inference_times = []
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
while (time.time() - start_time) < 10:
|
|||
|
|
ret, frame = cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# Resize到640x640
|
|||
|
|
frame = cv2.resize(frame, (640, 640))
|
|||
|
|
|
|||
|
|
# 推理
|
|||
|
|
infer_start = time.time()
|
|||
|
|
results = model.predict(
|
|||
|
|
frame,
|
|||
|
|
imgsz=640,
|
|||
|
|
conf=0.45,
|
|||
|
|
verbose=False,
|
|||
|
|
device=0, # GPU 0
|
|||
|
|
classes=[0] # person only
|
|||
|
|
)
|
|||
|
|
infer_end = time.time()
|
|||
|
|
|
|||
|
|
inference_times.append((infer_end - infer_start) * 1000)
|
|||
|
|
frame_count += 1
|
|||
|
|
|
|||
|
|
# 显示进度
|
|||
|
|
if frame_count % 10 == 0:
|
|||
|
|
elapsed = time.time() - start_time
|
|||
|
|
fps = frame_count / elapsed
|
|||
|
|
avg_latency = np.mean(inference_times)
|
|||
|
|
print(f" 帧数: {frame_count} | FPS: {fps:.1f} | 延迟: {avg_latency:.1f}ms")
|
|||
|
|
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print("\n⏹️ 测试被中断")
|
|||
|
|
|
|||
|
|
finally:
|
|||
|
|
cap.release()
|
|||
|
|
|
|||
|
|
# 统计结果
|
|||
|
|
elapsed = time.time() - start_time
|
|||
|
|
avg_fps = frame_count / elapsed
|
|||
|
|
|
|||
|
|
print(f"\n{'='*60}")
|
|||
|
|
print("测试结果")
|
|||
|
|
print(f"{'='*60}")
|
|||
|
|
print(f"总帧数: {frame_count}")
|
|||
|
|
print(f"测试时长: {elapsed:.1f}秒")
|
|||
|
|
print(f"平均FPS: {avg_fps:.1f}")
|
|||
|
|
print(f"平均推理延迟: {np.mean(inference_times):.1f}ms")
|
|||
|
|
print(f"P95推理延迟: {np.percentile(inference_times, 95):.1f}ms")
|
|||
|
|
print(f"P99推理延迟: {np.percentile(inference_times, 99):.1f}ms")
|
|||
|
|
print(f"{'='*60}")
|
|||
|
|
|
|||
|
|
print("\n✅ 测试完成！TensorRT 引擎工作正常")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
try:
|
|||
|
|
test_tensorrt_inference()
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"\n❌ 测试失败: {e}")
|
|||
|
|
import traceback
|
|||
|
|
traceback.print_exc()
|