import os
import cv2
import gradio as gr
import numpy as np
# from configs.config import Configuration
from scripts.charExtraction import CharExtraction
from scripts.bboxAnnotator import BBOXAnnotator

wPathPlat = "licence_plat.pt"
wPathChar = "licence_character.pt"
classList = np.array(['A','B','C','D','E','F','G','H','I','J','K','L','M',
                            'N','O','P','Q','R','S','T','U','V','W','X','Y','Z',
                            '0','1','2','3','4','5','6','7','8','9'])
sizePlat = (416,200)
# imageSample = Configuration['imageSample']


extractor = CharExtraction(wPlatePath=wPathPlat, wCharacterPath=wPathChar, classList=classList,
                            sizePlate=sizePlat, conf=0.5)
annotator = BBOXAnnotator()

def getAnnotatedImage(frame, conf):
    bbox, plateNum, confidence = extractor.predict(image=frame, conf=conf)
    annotateImg, plateNum = annotator.draw_bbox(frame, bbox, plateNum)
    prob_dict = {key: value for key, value in zip(plateNum, confidence)}

    # pathOu = os.path.join(os.path.abspath(''), "output/sample.jpg")
    # cv2.imwrite(pathOu, annotateImg)

    return annotateImg, prob_dict

def video_identity(frame, conf=0.45):
    image, plateNumber = getAnnotatedImage(frame, conf)

    return image, plateNumber


demo = gr.Interface(fn=video_identity,
                    inputs=[gr.Image(label='Input image'), gr.Slider(minimum=0.1, maximum=1.0, label='Conf value')],
                    outputs=[gr.Image(label='Result'), gr.Label(label='Plate Number Detected')],
                    examples = [["sample_image2.jpg"], ["sample_image3.jpg"], ["sample_image5.jpg"], ["sample_image6.jpg"]],
                    cache_examples=True)

if __name__ == "__main__":
    demo.queue()
    demo.launch()