import cv2
import time
import numpy as np
from rknnpool import rknnPoolExecutor
# 图像处理函数，实际应用过程中需要自行修改
from func import myFunc

cap = cv2.VideoCapture('/dev/video10')
# cap = cv2.VideoCapture(0)
modelPath = "./rknnModel/yolov5s_relu_tk2_RK3588_i8.rknn"
# 线程数, 增大可提高帧率
TPEs = 3
# 初始化rknn池
pool = rknnPoolExecutor(
    rknnModel=modelPath,
    TPEs=TPEs,
    func=myFunc)

# 初始化异步所需要的帧
if (cap.isOpened()):
    for i in range(TPEs + 1):
        ret, frame = cap.read()
        # 检查帧数据是否需要重塑
        if frame.size == 1280 * 720 * 3:  # 检查是否为扁平化数据
            frame = frame.reshape((720, 1280, 3)).astype(np.uint8)
        else:
            print(f"Unexpected frame shape: {frame.shape}")
        if not ret:
            cap.release()
            del pool
            exit(-1)
        pool.put(frame)

frames, loopTime, initTime = 0, time.time(), time.time()
while (cap.isOpened()):
    frames += 1
    ret, frame = cap.read()
    # 检查帧数据是否需要重塑
    if frame.size == 1280 * 720 * 3:  # 检查是否为扁平化数据
        frame = frame.reshape((720, 1280, 3)).astype(np.uint8)
    else:
        print(f"Unexpected frame shape: {frame.shape}")
    if not ret:
        break
    pool.put(frame)
    frame, flag = pool.get()
    if flag == False:
        break
    cv2.imshow('test', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    if frames % 30 == 0:
        print("30帧平均帧率:\t", 30 / (time.time() - loopTime), "帧")
        loopTime = time.time()

print("总平均帧率\t", frames / (time.time() - initTime))
# 释放cap和rknn线程池
cap.release()
cv2.destroyAllWindows()
pool.release()