Update app.py

app.py

@@ -15,38 +15,44 @@ 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)
 
 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
-    onnx_model = onnx.load(model_file)
-    onnx.checker.check_model(onnx_model)
-    ort_session = rt.InferenceSession(model_file)
-    return transform, ort_session 
+    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=model_file)
+# Инициализация модели
+transform, ort_session = initialize_model(model_file)
 
 # Создаем FastAPI приложение
 app = FastAPI()
 
 # Функция для получения текста
 def get_text(img_org):
-    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]
+    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")
@@ -55,10 +61,10 @@ async def predict(file: UploadFile = File(...)):
         # Получаем изображение из запроса
         image_bytes = await file.read()
         img = Image.open(BytesIO(image_bytes))
-        
+
         # Получаем текст с изображения
         result = get_text(img)
-        
+
         # Возвращаем распознанный текст
         return JSONResponse(content={"text": result})
     except Exception as e: