from functools import lru_cache

import gradio as gr
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

from cap import Predictor


@lru_cache()
def load_predictor(model):
    predictor = Predictor(hf_hub_download(
        f'7eu7d7/CAPTCHA_recognize',
        model,
    ), ckpt_name=model)
    return predictor


def process_image(image, model_name):
    """
    Process the uploaded image with selected model
    """
    if image is None:
        return "Please upload an image first"

    # Convert image to PIL format if needed
    if isinstance(image, np.ndarray):
        img = Image.fromarray(image.astype('uint8')).convert('RGB')
    else:
        img = image.convert('RGB')

    try:
        predictor = load_predictor(model_name)
        text = predictor.pred_img(img, show=False)
        return text
    except Exception as e:
        return f"Error processing image: {str(e)}"


# Create Gradio interface
with gr.Blocks(title="CAPTCHA Recognize") as demo:
    with gr.Row():
        # Left column - Input area
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="Upload CAPTCHA Image",
                type="pil",
                height=300
            )

            # Model selection dropdown
            model_dropdown = gr.Dropdown(
                label="Select Model",
                choices=[
                    "captcha-2000.safetensors",
                    "captcha-7400.safetensors",
                    "captcha-caformer-v2-6200.safetensors",
                    "captcha-caformer-v2-13000.safetensors",
                ],
                value="captcha-caformer-v2-13000.safetensors",  # 默认选择
                interactive=True
            )

            # Run button
            process_btn = gr.Button(
                "Run",
                variant="primary",
                size="lg"
            )

        # Right column - Output area
        with gr.Column(scale=1):
            text_output = gr.Textbox(
                label="Result",
                lines=4,
                interactive=False
            )

    # Bind events
    process_btn.click(
        fn=process_image,
        inputs=[image_input, model_dropdown],
        outputs=[text_output]
    )

# Launch the application
if __name__ == "__main__":
    demo.launch()