#!/usr/bin/env python
#version    :       2023.12.31
#language   :       ch
from maix import camera
from maix import display

from maix import image

from maix import nn
from maix.nn import decoder

import os
ScreenOrientation = False
try:
    if os.path.exists("/etc/cameraSize.cfg"):
        cameraSize = True
    else:
        cameraSize = False
except:
    cameraSize = False
def getLcdRotation(cameraCapture):
    global cameraSize
    if cameraSize:
        return lcdRotationNew(cameraCapture)
    else:
        return lcdRotation(cameraCapture)

def lcdRotationNew(inputImg):
    global cameraSize,ScreenOrientation
    imageRotationBuffer = inputImg.crop(0, 0, 320, 240)
    if ScreenOrientation:
        imgRotationAim = image.new(size = (240, 320))
        rotationAngle = 90
        GETROTATION = imageRotationBuffer.rotate(+rotationAngle, adjust=1)
    else:
        imgRotationAim = image.new(size = (320, 240))
        GETROTATION = imageRotationBuffer

    GETROTATION = imgRotationAim.draw_image(GETROTATION,0,0,alpha=1)
    return GETROTATION

def lcdRotation(inputImg):
    global cameraSize,ScreenOrientation
    imageRotationBuffer = inputImg.crop(0, 0, 240, 320)
    if ScreenOrientation:
        imgRotationAim = image.new(size = (240, 320))
        rotationAngle = 180
    else:
        imgRotationAim = image.new(size = (320, 240))
        rotationAngle = 90
    GETROTATION = imageRotationBuffer.rotate(+rotationAngle, adjust=1)
    GETROTATION = imgRotationAim.draw_image(GETROTATION,0,0,alpha=1)
    return GETROTATION
   
if cameraSize==True:
    camera.camera.config(size=(320,240))
else:
    camera.camera.config(size=(240,320))

image.load_freetype("/root/preset/fonts/SourceHanSansCN-Regular.otf")


class Number_recognition:
    mdsc_path = "/root/preset/model/Number.mud"
    labels = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
    anchors = [1.0, 5.0, 1.35, 5.42, 0.49, 2.55, 0.86, 3.75, 0.65, 4.38]

    def __init__(self):
        self.model = nn.load(self.mdsc_path)
        self.decoder = decoder.Yolo2(len(self.labels) , self.anchors , net_in_size = (224, 224) ,net_out_size = (7,7))
    def __del__(self):
        del self.model
        del self.decoder
    def cal_fps(self ,start , end):
        one_second = 1
        one_flash = end - start
        fps = one_second / one_flash
        return  fps

number_recognition = Number_recognition()
canvasImg = image.new(size = (240, 320))
while True:
    canvasImg.clear()
    img_mnist = getLcdRotation(camera.capture())
    img_mnist = img_mnist.crop(0, 0,224, 224)
    out = number_recognition.model.forward(img_mnist, quantize=1, layout = "hwc")
    boxes, probs = number_recognition.decoder.run(out, nms=0.5, threshold=0.3, img_size=(224,224))
    if len(boxes):
        for boxesi, box in enumerate(boxes):
            boxes[boxesi].append(probs[boxesi])
    if len(boxes):
        for i in (boxes):
            img_mnist.draw_string(i[0],i[1], (str(number_recognition.labels[i[4][0]]) + str(str(":") + str(round(i[4][1][i[4][0]]*100, 2)))), scale = 1, color = (255,0,0) , thickness = 1)
            img_mnist.draw_rectangle(i[0],i[1], int(i[0] + i[2]),int(i[1] + i[3]), color=(255,0,0), thickness=1)
    if ScreenOrientation:
        img_mnistVER = img_mnist.crop(0,0,240,320)
        img_mnistVER = img_mnistVER.rotate(-90, adjust=1)
        display.show(img_mnistVER)
    else:
        canvasImg.draw_image(img_mnist,48,8)
        canvasImg.draw_image((image.open("/root/preset/img/exit_ff0000_24x24.png")).rotate(0, adjust=0),288,216,alpha=1)
        display.show(canvasImg)