import gradio as gr
# import cv2
import os
import base64
from pathlib import Path
    
from core.poker_detector import PokerDetector

detector = PokerDetector(
    model_path="onnx/poker_detection_v4_rank.onnx"
)


### 构建 examples 

def build_examples():
    examples = []
    # 读取 examples 目录下的所有图片
    for file in os.listdir("examples"):
        if file.endswith(".jpg") or file.endswith(".png"):
            image_path = os.path.join("examples", file)
            examples.append([image_path])

    return examples


full_examples = build_examples()


with gr.Blocks(css="""
        .image img {
            max-height: 512px;
        }
    """
) as demo:
    gr.Markdown("""
                ## 扑克牌检测
                """
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="上传扑克牌图片", type="numpy", elem_classes="image")

        with gr.Column():
            with gr.Column():   
                result_image = gr.Image(
                    label="检测结果",
                    interactive=False,
                    visible=True,
                    elem_classes="image"
                )

            with gr.Column():
                use_time = gr.Textbox(
                    label="用时",
                    interactive=False,
                    visible=True,
                )
                
    with gr.Row():
        with gr.Column():
            gr.Examples(
                full_examples[:10], inputs=[image_input], label="示例图片",  examples_per_page=10,)
            

    def detect_poker(image):
        if image is None:
            return None, ""
        
        try:
            image_rgb_with_pred, time_info = detector.pred_and_draw(image)  

        except Exception as e:
            gr.Warning(f"检测失败: {e}")
            return None, "检测失败"


        return image_rgb_with_pred, time_info

    image_input.change(fn=detect_poker, 
                       inputs=[image_input], 
                       outputs=[result_image, use_time])

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