正点原子RK3588（二）——lenet测试和modelzoo模型

一、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:
                break
            if value > 0:
                topi = '{}: {}\n'.format(index[j], value)
            else:
                topi = '-1: 0.0\n'
            top5_str += topi
    print(top5_str)

if __name__ == '__main__':
    rknn_model = RKNN_MODEL
    rknn_lite = RKNNLite()
    # load RKNN model
    print('--> 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 environment
    print('--> Init runtime environment')
    ret = rknn_lite.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    # Inference
    print('--> 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 = 100


def process_video_yolov5(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件 {video_path}")
        return

    fps = int(cap.get(cv2.CAP_PROP_FPS))  # 获取视频帧率
    frame_count = 0
    vehicle_count_per_5_seconds = []
    current_5s_vehicle_count = 0
    total_vehicle_count = 0

    speed_per_5_seconds = []
    density_per_5_seconds = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 使用YOLOv5进行检测
        results = model(frame)

        # 解析检测结果，results.pandas().xyxy 返回检测结果的 DataFrame
        detections = 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 - x1
                speed = (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_length
            density_per_5_seconds.append(density)

            # 车速计算（当前帧车辆速度的平均值）
            avg_speed = np.mean(current_frame_speeds) if current_frame_speeds else 0
            speed_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'):
            break

    cap.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秒的时间点

    # 创建DataFrame
    df = 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')  # Windows
    save_dir = desktop_path

    if 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)

测试完毕