add FastAPI

app.py

@@ -4,71 +4,65 @@ 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
-
-
+from fastapi import FastAPI, File, UploadFile
+from io import BytesIO
+from fastapi.responses import JSONResponse
 
+# Инициализация модели
 model_file = "captcha.onnx"
 img_size = (32,128)
 charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
 tokenizer_base = Tokenizer(charset)
 
 def get_transform(img_size):
-        transforms = []
-        transforms.extend([
-            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):
     transform = get_transform(img_size)
-    # Onnx model loading
+    # Загрузка модели ONNX
     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)
+
+# Создаем FastAPI приложение
+app = FastAPI()
 
+# Функция для получения текста
 def get_text(img_org):
-    # 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)
-
-    # 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
-
-transform,ort_session = initialize_model(model_file=model_file)
-
-# Создание интерфейса
-with gr.Blocks() as demo:
-    image_input = gr.Image(type="pil")
-    text_output = gr.Textbox()
-    
-    # Кнопка для обработки изображения
-    submit_button = gr.Button("Распознать текст")
-    
-    # Связываем функцию с кнопкой
-    submit_button.click(fn=get_text, inputs=image_input, outputs=text_output)
-
-# Запуск с включенной поддержкой очереди
-demo.queue().launch()
+    preds, _ = tokenizer_base.decode(probs)
+    return preds[0]
 
-# if __name__ == "__main__":
-#     image_path = "8000.png"
-#     preds,probs = get_text(image_path)
-#     print(preds[0])
-     
+# Маршрут для обработки POST-запросов с изображениями
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    try:
+        # Получаем изображение из запроса
+        image_bytes = await file.read()
+        img = Image.open(BytesIO(image_bytes))
+        
+        # Получаем текст с изображения
+        result = get_text(img)
+        
+        # Возвращаем распознанный текст
+        return JSONResponse(content={"text": result})
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"message": str(e)})
 
+# Для запуска FastAPI приложения
+# uvicorn main:app --reload