app.py

import torch
import onnx
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

# Инициализация модели
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)

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()

# Функция для получения текста
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}")

# Маршрут для обработки 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