Update app.py

app.py

@@ -1,13 +1,22 @@
 import gradio as gr
 import torch
 import pickle
-from model import predict_with_beam_search, PosFormerImprovedWithIAC, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN, Vocab,IMG_WIDTH,IMG_HEIGHT
+from PIL import Image, ImageOps # <-- Добавляем импорт для работы с изображениями
 
-# --- загрузка модели (как в предыдущем примере) ---
+# Убедись, что model.py и vocab.pkl находятся в той же папке
+from model import predict_with_beam_search, PosFormerImprovedWithIAC, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN, Vocab, IMG_WIDTH, IMG_HEIGHT
+
+# --- 1. Загрузка модели (твой код здесь был в порядке) ---
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-with open("vocab.pkl", "rb") as f:
-    vocab = pickle.load(f)
 
+try:
+    with open("vocab.pkl", "rb") as f:
+        vocab = pickle.load(f)
+except FileNotFoundError:
+    print("Ошибка: Файл vocab.pkl не найден. Убедись, что он находится в той же директории.")
+    exit()
+
+# Инициализация модели
 model = PosFormerImprovedWithIAC(
     vocab=vocab,
     vocab_size=len(vocab.itos),
@@ -16,37 +25,64 @@ model = PosFormerImprovedWithIAC(
     id_to_token_map={i:s for s,i in vocab.stoi.items()}
 ).to(DEVICE)
 
-checkpoint = torch.load("d256 layers=4 densenet-121 32%.pth", map_location=DEVICE)
-model.load_state_dict(checkpoint["model_state_dict"])
-model.eval()
+try:
+    # Укажи правильный путь к своему файлу с весами
+    checkpoint_path = "d256 layers=4 densenet-121 32%.pth"
+    checkpoint = torch.load(checkpoint_path, map_location=DEVICE)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+    print(f"Модель и веса из '{checkpoint_path}' успешно загружены.")
+except FileNotFoundError:
+    print(f"Ошибка: Файл с весами '{checkpoint_path}' не найден.")
+    exit()
+except KeyError:
+    print("Ошибка: Ключ 'model_state_dict' не найден в файле с весами. Проверь структуру файла.")
+    exit()
+
+
+# --- 2. Функция-обертка для Gradio с предобработкой изображения ---
+def gradio_inference(image, beam_size=5):
+    """
+    Принимает "сырое" изображение от Gradio, предобрабатывает его и передает в модель.
+    """
+    if image is None:
+        return "Сначала нарисуйте что-нибудь на холсте."
+
+    # Предобработка изображения
+    # 1. Конвертируем в оттенки серого (L - Luminance)
+    processed_image = image.convert("L")
+
+    # 2. Инвертируем цвета. На холсте мы рисуем черным по белому.
+    #    Модели часто обучаются на белых символах на черном фоне.
+    processed_image = ImageOps.invert(processed_image)
+
+    # 3. Изменяем размер до того, который ожидает модель
+    #    Используем константы, импортированные из твоего файла model.py
+    processed_image = processed_image.resize((IMG_WIDTH, IMG_HEIGHT))
 
-# --- функция инференса ---
-def inference(image, beam_size=5):
-    # image здесь — PIL.Image из Sketchpad
-    return predict_with_beam_search(model, image, vocab, beam_size, DEVICE)
+    # Теперь передаем подготовленное изображение в твою функцию предсказания
+    return predict_with_beam_search(model, processed_image, vocab, beam_size, DEVICE)
 
-# --- интерфейс с полем рисования ---
+# --- 3. Интерфейс с правильным компонентом и функцией ---
 iface = gr.Interface(
-    fn=inference,
+    fn=gradio_inference, # <-- Используем нашу новую функцию-обертку
     inputs=[
-        gr.Image(
-            source="canvas",
-            tool="sketch",
-            type="pil",
+        gr.Sketchpad(
+            type="pil", # <-- Важно: получаем объект PIL.Image
             label="Нарисуйте формулу"
         ),
         gr.Slider(
             minimum=1,
             maximum=20,
-            value=10,        # <-- тут вместо default
+            value=10,
             step=1,
             label="Beam Size"
         )
     ],
     outputs=gr.Textbox(label="Предсказанная LaTeX-строка"),
     title="Handwritten Formula Recognition",
-    description="Нарисуйте формулу на холсте и получите LaTeX."
+    description="Нарисуйте формулу на холсте, чтобы получить ее в формате LaTeX. Модел�� лучше всего работает с одиночными, четко написанными символами и короткими формулами."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()