#!/usr/bin/env python
#version    :       2023.12.31
#language   :       en
from time import time
from maix import image

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
   

image.load_freetype("/root/preset/fonts/SourceHanSansCN-Regular.otf")
class Hand:
    mud_path = "/root/preset/model/hand_int8.mud"
    labels = ["0","1","2","3","4","5"]
    anchors = [3.78, 5.81, 3.97, 3.98, 4.05, 4.98, 4.81, 5.41, 2.91, 4.53]

    def __init__(self) -> None:
        from maix import nn
        self.model = nn.load(self.mud_path)
        from maix.nn import decoder
        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

    def draw_rectangle_with_title(self ,img, box, disp_str , fps ):
        img.draw_rectangle(box[0], box[1], box[0] + box[2], box[1] + box[3],color=(255, 0, 0), thickness=2)
        img.draw_string(box[0], box[1]+ box[3] ,disp_str, scale=1,color=(0, 0, 255), thickness=2)
        img.draw_string(0, 0 ,'FPS :'+str(fps), scale=2 ,color=(0, 0, 255), thickness=2)

    def process(self,input):
        t =  time()
        out = self.model.forward(input, quantize=1, layout = "hwc")
        boxes, probs = self.decoder.run(out, nms=0.5, threshold=0.6, img_size=(224,224))
        for i, box in enumerate(boxes):
            class_id = probs[i][0]
            prob = probs[i][1][class_id]
            disp_str = "{}:{:.2f}%".format(self.labels[class_id], prob*100)
            fps = self.cal_fps(t, time())
            self.draw_rectangle_with_title(input, box, disp_str, fps)


def main():
    from maix import display, camera,image
    global Hand
    app =  Hand()
    canvasImg = image.new(size = (320, 240))
    while True:
        img = getLcdRotation(camera.capture()).crop(0, 0,224, 224)
        app.process(img)
        canvasImg.draw_image(img,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)

main()