# Untitled - By: nian0 - 周一 11月 20 2023

import sensor, image, lcd, time
from fpioa_manager import fm
import KPU as kpu
import gc
from Maix import GPIO
from machine import Timer, PWM, UART, Timer

# 摄像头初始化
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_vflip(1)  # 摄像头后置方式
sensor.set_pixformat(sensor.RGB565)     # 设置像素格式为彩色 RGB565
#sensor.set_pixformat(sensor.GRAYSCALE)  # 设置像素格式为灰色
sensor.set_framesize(sensor.QVGA)       # 设置帧大小为 QVGA (320x240)
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(False)
sensor.set_auto_gain(0,0)

# 初始化 LCD
lcd.init()

# 从SD卡加载模型
# task = kpu.load("/sd/model-96249.kmodel")
labels = ['green', 'light', 'yellow', 'red']
anchors = [0.88, 3.47, 0.69, 0.81, 0.97, 3.47, 1.12, 3.75, 1.25, 4.19]



while True:
    task = None
    task = kpu.load("/sd/model-96249.kmodel")
    # 初始化yolo2网络，识别可信概率为0.7（70%）
    a = kpu.init_yolo2(task, 0.7, 0.3, 5, anchors)
    img = sensor.snapshot()

    # 调整图像大小
    img = img.resize(224, 224)

    # # 在LCD屏幕上显示图像
    # lcd.display(img)

    # 准备图像进行推理
    img.pix_to_ai()

    code = kpu.run_yolo2(task, img)  # 运行yolo2网络

    print(code)
    if code:
        for obj in code:
            pos = obj.rect()
            print(labels[obj.classid()])
            img.draw_rectangle(pos)
            img.draw_string(pos[0], pos[1], "%s : %.2f" %(labels[obj.classid()], obj.value()), scale=2, color=(255, 0, 0))

        #comm.send_detect_result(code, classes)
            #img.draw_string(0, 200, "t:%dms" %(t), scale=2, color=(255, 0, 0))
            a = lcd.display(img)
    else:
        a = lcd.display(img)


    if not task is None:
        kpu.deinit(task)