import numpy as np
import gradio as gr
from PIL import Image
from grounded_sam_plate import recognize_plate

CUT_RATIO = 0.15

def process(input_img, cut_ratio):
    result = recognize_plate(input_img, cut_ratio=cut_ratio)
    return result['plate'], result['rectified_plate'], result['detection'], result['chinese']


with gr.Blocks() as demo:
    gr.Markdown(
        """
        please visit https://github.com/ZhengJun-AI/license_plate_recognition.

        this demo only support some Chinese characters: [鄂, 桂, 黑, 冀, 贵, 京, 鲁, 闽, 苏, 皖, 豫, 粤, 新, 川, 吉, 津, 辽, 陕, 浙, 蒙]
        """
    )

    with gr.Row():
        image_input = gr.Image(label="Input Image")
        image_output = gr.Image(label="Rectified Plate", height=300)
    with gr.Accordion("Check more details", open=False):
        with gr.Row():
            cut_ratio = gr.Slider(0, 1, value=CUT_RATIO, step=0.01, label="Cut ratio")
            chinese = gr.Image(label="Chinese", height=100)
        image_detection = gr.Image(label="Detection result")
    result = gr.Textbox(lines=1, label="Result")
    image_button = gr.Button("Run")
    image_button.click(process, inputs=[image_input, cut_ratio], outputs=[result, image_output, image_detection, chinese])

    gr.Markdown("## Examples")
    gr.Examples(
        examples=[
            ['./images/0.jpg', 0.15],
            ['./images/1.jpeg', 0.15],
            ['./images/6.jpeg', 0.15],
            ['./images/5.jpg', 0.15],
        ],
        inputs=[image_input, cut_ratio],
        outputs=[result, image_output, image_detection, chinese],
        fn=process,
        cache_examples=True,
    )


demo.launch()