from media.sensor import *
from media.display import *
from media.media import *
from libs.YOLO import YOLOv5
from libs.Utils import *
import os,sys,gc
import ulab.numpy as np
import image

if __name__=="__main__":
    # 这里仅为示例，自定义场景请修改为您自己的模型路径、标签名称、模型输入大小
    kmodel_path="/data/gangqiu_det.kmodel"
    labels = ["gangqiu"]
    model_input_size=[224,224]
    # 添加显示模式，默认hdmi，可选hdmi/lcd/lt9611/st7701/hx8399/nt35516/nt35532/gc9503/aml020t/jd9852/ili9806/virt；其中hdmi默认对应lt9611，lcd默认对应st7701
    display_mode="lcd"
    # 显示分辨率，None表示使用当前显示屏默认分辨率；使用virt时可在这里手动设置，例如[800, 480]
    display_size=[800,480]
    rgb888p_size=[224,224]
    confidence_threshold = 0.5
    nms_threshold=0.45
    count=0

    sensor = Sensor(id=2)
    sensor.reset()
    #sensor.set_hmirror(False)
    #sensor.set_vflip(False)
    sensor.set_framesize(width = display_size[0], height = display_size[1],chn=CAM_CHN_ID_0)
    sensor.set_pixformat(Sensor.YUV420SP,chn=CAM_CHN_ID_0)
    sensor.set_framesize(width = rgb888p_size[0], height = rgb888p_size[1], chn=CAM_CHN_ID_1)
    sensor.set_pixformat(Sensor.RGBP888, chn=CAM_CHN_ID_1)
    sensor_bind_info = sensor.bind_info(x = 0, y = 0, chn = CAM_CHN_ID_0)
    Display.bind_layer(**sensor_bind_info, layer = Display.LAYER_VIDEO1)
    osd_img = image.Image(display_size[0], display_size[1], image.ARGB8888)
    Display.init(Display.ST7701,width=display_size[0],height=display_size[1],osd_num=1, to_ide = True)
    sensor.run()

    # 初始化YOLOv5实例
    yolo=YOLOv5(task_type="detect",mode="video",kmodel_path=kmodel_path,labels=labels,rgb888p_size=rgb888p_size,model_input_size=model_input_size,display_size=display_size,conf_thresh=confidence_threshold,nms_thresh=nms_threshold,max_boxes_num=50,debug_mode=0)
    yolo.config_preprocess()
    c=utime.clock()
    while True:
        c.tick()
        count+=1
        img=sensor.snapshot(chn=CAM_CHN_ID_1)
        img_np=img.to_numpy_ref()
        res=yolo.run(img_np)
        yolo.draw_result(res,osd_img)
        Display.show_image(osd_img)
        if count==10:
            gc.collect()
            count=0
        print("FPS:",c.fps())
    yolo.deinit()
    sensor.stop()
    Display.deinit()
    nn.shrink_memory_pool()