import sentry, sensor, image, time, lcd, gc, cmath, os
from maix import KPU

uart = sentry.UART(115200)

folder = os.listdir('/sd')
if 'image' not in folder:
    os.mkdir('/sd/image')

lcd.init()                          # Init lcd display
lcd.clear(lcd.RED)                  # Clear lcd screen.

# sensor.reset(dual_buff=True)      # improve fps
sensor.reset()                      # Reset and initialize the sensor.
sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QVGA)   # Set frame size to QVGA (320x240)
sensor.set_vflip(True)              # Flip the cam
sensor.set_hmirror(True)            # Mirror the cam
sensor.skip_frames(time = 1000)     # Wait for settings take effect.
clock = time.clock()                # Create a clock object to track the FPS.

print("ready load model")

labels = ["polyp"] # object type names, follow the order in "abel.txt
anchor = (1.67, 0.80, 1.61, 1.17, 1.92, 1.47, 3.08, 2.28, 4.42, 3.11) # anchors, use the value from the second line of anchor.txt

kpu = KPU()
# 从sd或flash加载模型
kpu.load_kmodel('/sd/det.kmodel')
# kpu.load_kmodel(0x300000, 584744)
kpu.init_yolo2(anchor, anchor_num=(int)(len(anchor)/2), img_w=320, img_h=240, net_w=320 , net_h=240 ,layer_w=10 ,layer_h=8, threshold=0.6, nms_value=0.3, classes=len(labels))

def image_save():
    img2 = sensor.snapshot()
    for i in range(3):
        path = "/sd/image/image_%d.jpg"%i
        img2.save(path)
        time.sleep(0.2)
        
uart.write('start')

while(True):
    gc.collect()

    clock.tick()
    img = sensor.snapshot()

    kpu.run_with_output(img)
    dect = kpu.regionlayer_yolo2()

    fps = clock.fps()

    if len(dect) > 0:
        for l in dect :
            a = img.draw_rectangle(l[0],l[1],l[2],l[3],color=(0,255,0))

            info = "%s %.3f" % (labels[l[4]], l[5])
            a = img.draw_string(l[0],l[1],info,color=(255,0,0),scale=2.0)
            print(info)
            del info
        uart.write('1')

    a = img.draw_string(0, 0, "%2.1ffps" %(fps),color=(0,60,255),scale=2.0)
    lcd.display(img)

    if uart.read():
        image_save()