Update app.py

app.py

@@ -4,71 +4,66 @@ import onnxruntime as rt
 from torchvision import transforms as T
 from PIL import Image
 from tokenizer_base import Tokenizer
-from fastapi import FastAPI, File, UploadFile
-from io import BytesIO
-from fastapi.responses import JSONResponse
+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)
+        transforms = []
+        transforms.extend([
+            T.Resize(img_size, T.InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(0.5, 0.5)
+        ])
+        return T.Compose(transforms)
 
 def to_numpy(tensor):
     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
 def initialize_model(model_file):
-    try:
-        # Загрузка модели ONNX
-        onnx_model = onnx.load(model_file)
-        onnx.checker.check_model(onnx_model)
-        ort_session = rt.InferenceSession(model_file)
-        transform = get_transform(img_size)
-        return transform, ort_session
-    except Exception as e:
-        raise RuntimeError(f"Ошибка при инициализации модели: {e}")
-
-# Инициализация модели
-transform, ort_session = initialize_model(model_file)
-
-# Создаем FastAPI приложение
-app = FastAPI()
+    transform = get_transform(img_size)
+    # Onnx model loading
+    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):
-    try:
-        x = transform(img_org.convert('RGB')).unsqueeze(0)
-        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, _ = tokenizer_base.decode(probs)
-        return preds[0]
-    except Exception as e:
-        raise RuntimeError(f"Ошибка при обработке изображения: {e}")
+    # img_org = Image.open(image_path)
+    # Preprocess. Model expects a batch of images with shape: (B, C, H, W)
+    x = transform(img_org.convert('RGB')).unsqueeze(0)
 
-# Маршрут для обработки POST-запросов с изображениями
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    try:
-        # Получаем изображение из запроса
-        image_bytes = await file.read()
-        img = Image.open(BytesIO(image_bytes))
+    # compute ONNX Runtime output prediction
+    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]
+    print(preds)
+    return preds
 
-        # Получаем текст с изображения
-        result = get_text(img)
+transform,ort_session = initialize_model(model_file=model_file)
 
-        # Возвращаем распознанный текст
-        return JSONResponse(content={"text": result})
-    except Exception as e:
-        return JSONResponse(status_code=500, content={"message": str(e)})
+gr.Interface(
+    get_text,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(),
+    title="Text Captcha Reader",
+    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"]
+).launch()
 
-# Для запуска FastAPI приложения
-# uvicorn main:app --reload
+# if __name__ == "__main__":
+#     image_path = "8000.png"
+#     preds,probs = get_text(image_path)
+#     print(preds[0])
+