Update app.py

app.py

@@ -4,11 +4,10 @@ import onnxruntime as rt
 from torchvision import transforms as T
 from PIL import Image
 from tokenizer_base import Tokenizer
-import io
-import base64
-from fastapi import FastAPI, UploadFile, File
-from pydantic import BaseModel
-import numpy as np
+import pathlib
+import os
+import gradio as gr
+from huggingface_hub import Repository
 
 # Параметры модели
 model_file = "captcha.onnx"
@@ -16,60 +15,54 @@ img_size = (32,128)
 charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
 tokenizer_base = Tokenizer(charset)
 
-# Инициализация преобразования изображений
+# Преобразования для изображения
 def get_transform(img_size):
-    transforms = []
-    transforms.extend([
+    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 
+    return transform, ort_session
 
-# Инициализация модели
-transform, ort_session = initialize_model(model_file=model_file)
-
-# Функция для получения текста с изображения
+# Функция для распознавания текста на изображении
 def get_text(img_org):
-    x = transform(img_org.convert('RGB')).unsqueeze(0)
+    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]
+    probs = torch.tensor(logits).softmax(-1)  # Вычисляем вероятности
+    preds, probs = tokenizer_base.decode(probs)  # Декодируем результат
+    preds = preds[0]  # Получаем первое предсказание
     return preds
 
-# FastAPI приложение
-app = FastAPI()
-
-# Модель данных для работы с запросом
-class ImageData(BaseModel):
-    image: str
-
-# Endpoint для обработки изображения в формате Base64
-@app.post("/predict/")
-async def predict(data: ImageData):
-    try:
-        # Декодируем Base64
-        img_data = base64.b64decode(data.image.split(",")[1])
-        img = Image.open(io.BytesIO(img_data))
-        result = get_text(img)
-        return {"result": result}
-    except Exception as e:
-        return {"error": str(e)}
+# Инициализация модели
+transform, ort_session = initialize_model(model_file=model_file)
 
-# Запуск сервера (для локального тестирования)
-# Если вы хотите запустить сервер на хосте, используйте команду:
-# uvicorn filename:app --reload
+# Создание 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()