# -*- coding: utf-8 -*-
# @Time    : 2024/7/6 下午12:23
# @Author  : len
# @File    : 红绿灯模型.py
# @Software: PyCharm 
# @Comment :

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

input_size = (224, 224)
labels = ['green', 'yellow', 'red']
anchors = [1.25, 1.19, 1.41, 1.31, 1.12, 1.09, 1.47, 1.47, 1.62, 1.5]
QR_num = 0
QR_Flag = False
JTD_num = 0
JTD_Flag = False
# model_addr="/sd/model-JTD.kmodel"
model_addr="0x500000"

#映射串口引脚
fm.register(6, fm.fpioa.UART1_RX, force=True)
fm.register(7, fm.fpioa.UART1_TX, force=True)

#初始化串口
uart = UART(UART.UART1, 115200, read_buf_len=4096)

DISTORTION_FACTOR = 1.5 # 设定畸变系数
def init_sensor():
    '''
    初始化感光芯片
    '''
    lcd.init(freq=15000000)             #初始化LCD
    sensor.reset()                      #复位和初始化摄像头，执行sensor.run(0)停止。
    sensor.set_vflip(1)                 #将摄像头设置成后置方式（所见即所得）
    sensor.set_pixformat(sensor.RGB565) # 设置像素格式为彩色 RGB565
    #sensor.set_pixformat(sensor.GRAYSCALE) # 设置像素格式为灰色
    sensor.set_framesize(sensor.QVGA)   # 设置帧大小为 QVGA (320x240)
    sensor.set_windowing(input_size)
    sensor.set_vflip(1)   # 打开垂直翻转 如果是 01Studio 的 K210 不开启会导致画面方向与运动方向相反
    sensor.set_hmirror(1)   # 打开水平镜像 如果是 01Studio 的 K210 不开启会导致画面方向与运动方向相反
    sensor.skip_frames(time = 2000)     # 等待设置生效.
    clock = time.clock()                # 创建一个时钟来追踪 FPS（每秒拍摄帧数）

init_sensor()

JTD_Flag = True

def JTD_Check():
    '''
    红绿灯函数
    '''
    global JTD_Flag
    list1 = bytearray([0x55,0x02,0x92,0x00,0x00,0x00,0x00,0xBB])
    if JTD_Flag:
        print("正在识别")
        try:
            task = None
            print("正在加载模型")
            task = kpu.load(model_addr)
            print("加载完成")
            kpu.init_yolo2(task, 0.85, 0.25, 5, anchors) # 初始化 YOLOv2 模型
            while JTD_Flag:
                img = sensor.snapshot()
                objects = kpu.run_yolo2(task, img)
                if objects:
                    for obj in objects:
                        pos = obj.rect()
                        img.draw_rectangle(pos)
                        img.draw_string(pos[0], pos[1], "%s : %.2f" %(labels[obj.classid()], obj.value()), scale=2, color=(255, 0, 0))
                        print( labels[obj.classid()], obj.value() )
                    if( labels[obj.classid()] == 'green'):
                        list1[3] = 0x06
                        uart.write(list1)
                        print(list1)
                        print("绿灯")
                        JTD_Flag = False
                    elif( labels[obj.classid()] == 'yellow' ):
                        list1[3] = 0x07
                        uart.write(list1)
                        print(list1)
                        print("黄灯")
                        JTD_Flag = False
                    elif( labels[obj.classid()] == 'red' ):
                        list1[3] = 0x05
                        uart.write(list1)
                        print(list1)
                        print("红灯")
                        JTD_Flag = False
                    else:
                        list1[3] = 0x05
                        uart.write(list1)
                        JTD_Flag = False
                lcd.display(img)
        except Exception as e:
            raise e
        finally:
            if not task is None:
                kpu.deinit(task)

while True:
    JTD_Flag = True
    # 交通灯识别
    JTD_Check()