import spaces
import sys
print(sys.version_info)

import gradio as gr
#from florence import model_init, draw_image
#from wikai import analyze_dial, ocr_and_od
import matplotlib.pyplot as plt
from aimodel import model_init, read_meter, visualization
from test_rect import read_meter as read_meter_rect
from PIL import Image
from io import BytesIO
import logging
#logging.basicConfig(level=logging.DEBUG)
print("Loading model...")
fl, fl_ft = model_init(hack=False)
@spaces.GPU(duration=30)
def process_image(input_image:Image, meter_type:str):
    if meter_type == "方形儀表":    
        value, img = read_meter_rect(input_image, fl, fl_ft)
        return img, f"辨識結果: PA={value}", None
    assert meter_type == "圓形儀表"    
    plt.clf()
    print("process_image")
    W, H = 640, 480
    if input_image is None:
        return None, None    
    meter_result  = read_meter(input_image, fl, fl_ft)
    img, fig = visualization(meter_result)
    buf = BytesIO()
    fig.savefig(buf, format='png')
    buf.seek(0)
    fig_img = Image.open(buf)
    plt.clf()
    return img, f"辨識結果: PSI={meter_result.needle_psi:.1f}  kg/cm²={meter_result.needle_kg_cm2:.2f} ", fig_img

with gr.Blocks() as demo:

    gr.Markdown("## 指針辨識系統\n請選擇儀表類型，上傳圖片，或點擊Submit")

    with gr.Row():
        with gr.Column():
            with gr.Row():
                clear_button = gr.ClearButton()
                submit_button = gr.Button("Submit", variant="primary")
                meter_type_dropdown = gr.Dropdown(choices=["圓形儀表", "方形儀表"], label="選擇選項")
            image_input = gr.Image(type="pil", label="上傳圖片")
        with gr.Column():
            number_output = gr.Textbox(label="辨識結果", placeholder="辨識結果")
            image_output = gr.Image(label="輸出圖片")
            plot_output = gr.Image(label="模型結果")

    clear_button.add([image_input, image_output, number_output])

    image_input.upload(
        fn=process_image,
        inputs=[image_input, meter_type_dropdown],
        outputs=[image_output, number_output, plot_output],
        queue=False
    )

    submit_button.click(
        fn=process_image,
        inputs=[image_input, meter_type_dropdown],
        outputs=[image_output, number_output, plot_output],
    )

demo.launch(debug=True)