文章目录
- 一、lenet
- 二、modelzoo模型
- 2.1 介绍
- 2.2 测试
一、lenet
import cv2
import numpy as np
from rknnlite.api import RKNNLite
RKNN_MODEL = 'LeNet5_mnist_model.rknn'def show_top5(result):output = result[0].reshape(-1)output_sorted = sorted(output, reverse=True)top5_str = 'LeNet\n-----TOP 5-----\n'for i in range(5):value = output_sorted[i]index = np.where(output == value)for j in range(len(index)):if (i + j) >= 5:breakif value > 0:topi = '{}: {}\n'.format(index[j], value)else:topi = '-1: 0.0\n'top5_str += topiprint(top5_str)if __name__ == '__main__':rknn_model = RKNN_MODELrknn_lite = RKNNLite()# load RKNN modelprint('--> Load RKNN model')ret = rknn_lite.load_rknn(rknn_model)if ret != 0:print('Load RKNN model failed')exit(ret)print('done')ori_img = cv2.imread('./5.png')img = cv2.cvtColor(ori_img, cv2.COLOR_RGB2GRAY)img = np.expand_dims(img,0)img = np.expand_dims(img,0)# init runtime environmentprint('--> Init runtime environment')ret = rknn_lite.init_runtime()if ret != 0:print('Init runtime environment failed')exit(ret)# Inferenceprint('--> Running model')outputs = rknn_lite.inference(inputs=[img])show_top5(outputs)print('done')rknn_lite.release()
二、modelzoo模型
2.1 介绍
modelzoo提供了大量模型转换与部署的示例代码,旨在帮助用户快速在搭载瑞芯微芯片的开发板上运行各种模型。
2.2 测试
刚刚运行了
python3 mobilenet.py --target rk3588
发现缺少numpy,运行conda install numpy
,安装即可
嗯,少个cv2,pip install opencv-python
、有少个rknn
感觉不对劲,我又把rknn-toolkit2给装了,装玩之后,运行python3 mobilenet.py --target rk3588
,运行成功。想了想我在windows底下安装的pytorch,算了,好像不需要这个toolkit。
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
import pandas as pd# 加载YOLOv5模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)# 定义观测区域的长度(单位:米),假设为100米
observation_length = 100def process_video_yolov5(video_path):cap = cv2.VideoCapture(video_path)if not cap.isOpened():print(f"无法打开视频文件 {video_path}")returnfps = int(cap.get(cv2.CAP_PROP_FPS)) # 获取视频帧率frame_count = 0vehicle_count_per_5_seconds = []current_5s_vehicle_count = 0total_vehicle_count = 0speed_per_5_seconds = []density_per_5_seconds = []while True:ret, frame = cap.read()if not ret:break# 使用YOLOv5进行检测results = model(frame)# 解析检测结果,results.pandas().xyxy 返回检测结果的 DataFramedetections = results.xyxy[0].cpu().numpy() # [x1, y1, x2, y2, conf, cls]current_frame_speeds = []for *box, conf, cls in detections:# 检测类别ID,2: 汽车, 3: 摩托车, 5: 公共汽车, 7: 卡车if int(cls) in [2, 3, 5, 7]:x1, y1, x2, y2 = map(int, box)current_5s_vehicle_count += 1 # 累计每帧中检测到的车辆total_vehicle_count += 1# 假设简单的速度估计(可以根据车辆框的大小或位移进一步优化)width = x2 - x1speed = (width / observation_length) * fps # 简单估计速度current_frame_speeds.append(speed)# 在图像上绘制检测框cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)frame_count += 1# 每5秒钟(fps * 5 帧),记录一次车辆流量和速度if frame_count % (fps * 5) == 0:vehicle_count_per_5_seconds.append(current_5s_vehicle_count)print(f"第{len(vehicle_count_per_5_seconds) * 5}秒,车辆流量: {current_5s_vehicle_count}")# 车流密度计算(每5秒的车辆数 / 观测区域长度)density = current_5s_vehicle_count / observation_lengthdensity_per_5_seconds.append(density)# 车速计算(当前帧车辆速度的平均值)avg_speed = np.mean(current_frame_speeds) if current_frame_speeds else 0speed_per_5_seconds.append(avg_speed)current_5s_vehicle_count = 0 # 重置当前5秒的车辆计数# 显示结果cv2.imshow('Vehicle Detection with YOLOv5', frame)# 按 'q' 键退出if cv2.waitKey(1) & 0xFF == ord('q'):breakcap.release()cv2.destroyAllWindows()print(f"视频 {video_path} 处理完成,车辆总数: {total_vehicle_count}")return vehicle_count_per_5_seconds, density_per_5_seconds, speed_per_5_seconds# 绘制并保存车流密度、流量、速度的折线图到桌面,并生成Excel文件
def save_to_excel_and_plot(vehicle_counts, density_counts, speed_counts):time_points = list(range(5, len(vehicle_counts) * 5 + 1, 5)) # X轴为5秒的时间点# 创建DataFramedf = pd.DataFrame({'Time (seconds)': time_points,'Vehicle Count (Flow)': vehicle_counts,'Vehicle Density': density_counts,'Average Speed (m/s)': speed_counts})# 获取桌面路径desktop_path = os.path.join(os.path.join(os.environ['USERPROFILE']), 'Desktop') # Windowssave_dir = desktop_pathif not os.path.exists(save_dir):os.makedirs(save_dir)# 保存Excel文件excel_path = os.path.join(save_dir, 'vehicle_data.xlsx')df.to_excel(excel_path, index=False)print(f"数据已保存到 {excel_path}")# 绘制折线图plt.figure(figsize=(10, 6))# 绘制车辆流量plt.plot(time_points, vehicle_counts, marker='o', linestyle='-', color='b', label='Vehicle Flow')# 绘制车流密度plt.plot(time_points, density_counts, marker='x', linestyle='--', color='r', label='Vehicle Density')# 绘制平均速度plt.plot(time_points, speed_counts, marker='s', linestyle='-', color='g', label='Average Speed')plt.title('Vehicle Data over Time')plt.xlabel('Time (seconds)')plt.ylabel('Value')plt.grid(True)plt.legend()# 保存图表到桌面save_path = os.path.join(save_dir, 'vehicle_data_plot.png')plt.savefig(save_path)print(f"图表已保存到 {save_path}")plt.show()# 处理视频
video_paths = ['20240501_20240501125647_20240501140806_125649.mp4']
for video_path in video_paths:vehicle_counts, density_counts, speed_counts = process_video_yolov5(video_path)save_to_excel_and_plot(vehicle_counts, density_counts, speed_counts)
测试完毕