import torch
import onnx
import onnxruntime as rt
from torchvision import transforms as T
from PIL import Image
from tokenizer_base import Tokenizer
import pathlib
import os
import gradio as gr
from huggingface_hub import Repository

# Параметры модели
model_file = "captcha.onnx"
img_size = (32,128)
charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
tokenizer_base = Tokenizer(charset)

# Преобразования для изображения
def get_transform(img_size):
    transforms = [
        T.Resize(img_size, T.InterpolationMode.BICUBIC),
        T.ToTensor(),
        T.Normalize(0.5, 0.5)
    ]
    return T.Compose(transforms)

# Преобразование тензора в numpy
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# Инициализация модели ONNX
def initialize_model(model_file):
    transform = get_transform(img_size)
    onnx_model = onnx.load(model_file)
    onnx.checker.check_model(onnx_model)
    ort_session = rt.InferenceSession(model_file)
    return transform, ort_session

# Функция для распознавания текста на изображении
def get_text(img_org):
    x = transform(img_org.convert('RGB')).unsqueeze(0)  # Преобразуем изображение в тензор

    # Предсказание с помощью ONNX
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
    logits = ort_session.run(None, ort_inputs)[0]
    probs = torch.tensor(logits).softmax(-1)  # Вычисляем вероятности
    preds, probs = tokenizer_base.decode(probs)  # Декодируем результат
    preds = preds[0]  # Получаем первое предсказание
    return preds

# Инициализация модели
transform, ort_session = initialize_model(model_file=model_file)

# Создание Gradio интерфейса для обработки изображений
def gradio_interface():
    return gr.Interface(
        fn=get_text,
        inputs=gr.Image(type="pil"),  # Принимаем изображение в формате PIL
        outputs=gr.Textbox(),  # Выводим текст
        title="Text Captcha Reader",  # Заголовок
        description="Распознавание текста на изображениях капчи.",  # Описание
        examples=["8000.png", "11JW29.png", "2a8486.jpg", "2nbcx.png",  # Примеры изображений
                  "000679.png", "000HU.png", "00Uga.png.jpg", "00bAQwhAZU.jpg",
                  "00h57kYf.jpg", "0EoHdtVb.png", "0JS21.png", "0p98z.png", "10010.png"]
    )

# Запуск интерфейса
gradio_interface().launch()