Update app.py

app.py

@@ -1,74 +1,87 @@
 import gradio as gr
 import torch
 import pickle
-from PIL import Image, ImageOps # <-- Добавляем импорт для работы с изображениями
+from PIL import Image, ImageOps
+import torchvision.transforms as transforms
 
 # Убедись, что 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. Загрузка модели (твой код здесь был в порядке) ---
+# --- 1. Загрузка модели (без изменений) ---
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# (Загрузка vocab и модели - твой код здесь был правильный, я его сокращу для краткости)
 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),
-    pad_token_id=vocab.stoi[PAD_TOKEN],
-    structure_symbols_set={PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN, "{","}","^","_"},
-    id_to_token_map={i:s for s,i in vocab.stoi.items()}
-).to(DEVICE)
+    model = PosFormerImprovedWithIAC(
+        vocab=vocab,
+        vocab_size=len(vocab.itos),
+        pad_token_id=vocab.stoi[PAD_TOKEN],
+        structure_symbols_set={PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN, "{","}","^","_"},
+        id_to_token_map={i:s for s,i in vocab.stoi.items()}
+    ).to(DEVICE)
 
-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' не найден в файле с весами. Проверь структуру файла.")
+    print("Модель успешно загружена.")
+except Exception as e:
+    print(f"Ошибка при загрузке модели или словаря: {e}")
     exit()
 
 
-# --- 2. Функция-обертка для Gradio с предобработкой изображения ---
-def gradio_inference(image, beam_size=5):
+# --- 2. Функция предобработки (PIL Image -> Tensor) ---
+def preprocess_image(image: Image.Image) -> torch.Tensor:
+    """Подготавливает PIL Image для модели."""
+    # Конвертируем в оттенки серого
+    img = image.convert("L")
+    # Инвертируем (белые символы на черном фоне)
+    img = ImageOps.invert(img)
+    # Изменяем размер
+    img = img.resize((IMG_WIDTH, IMG_HEIGHT))
+
+    # Пайплайн для конвертации в тензор и нормализации [0, 1]
+    transform_pipeline = transforms.Compose([
+        transforms.ToTensor(),
+        # transforms.Normalize((0.5,), (0.5,)) # Раскомментируй, если модель этого требует
+    ])
+
+    tensor = transform_pipeline(img)
+    # Добавляем "батч" измерение (B, C, H, W) -> [1, 1, H, W]
+    tensor = tensor.unsqueeze(0)
+
+    return tensor
+
+
+# --- 3. Функция-обертка для Gradio (ИСПРАВЛЕННАЯ) ---
+def gradio_inference(sketchpad_data, beam_size=5):
     """
-    Принимает "сырое" изображение от Gradio, предобрабатывает его и передает в модель.
+    Принимает данные от Sketchpad (словарь), извлекает изображение,
+    предобрабатывает его и передает в модель.
     """
-    if image is None:
+    if sketchpad_data is None or sketchpad_data['composite'] is None:
         return "Сначала нарисуйте что-нибудь на холсте."
 
-    # Предобработка изображения
-    # 1. Конвертируем в оттенки серого (L - Luminance)
-    processed_image = image.convert("L")
-
-    # 2. Инвертируем цвета. На холсте мы рисуем черным по белому.
-    #    Модели часто обучаются на белых символах на черном фоне.
-    processed_image = ImageOps.invert(processed_image)
+    # !!! ГЛАВНОЕ ИСПРАВЛЕНИЕ ЗДЕСЬ !!!
+    # Извлекаем итоговое изображение (PIL Image) из словаря, который прислал Gradio
+    pil_image = sketchpad_data['composite']
 
-    # 3. Изменяем размер до того, который ожидает модель
-    #    Используем константы, импортированные из твоего файла model.py
-    processed_image = processed_image.resize((IMG_WIDTH, IMG_HEIGHT))
+    # Предобрабатываем PIL Image в тензор
+    image_tensor = preprocess_image(pil_image)
 
-    # Теперь передаем подготовленное изображение в твою функцию предсказания
-    return predict_with_beam_search(model, processed_image, vocab, beam_size, DEVICE)
+    # Выполняем предсказание (передаем тензор)
+    # Убедись, что твоя функция predict_with_beam_search ожидает тензор первым аргументом для изображения
+    return predict_with_beam_search(model, image_tensor, vocab, beam_size, DEVICE)
 
-# --- 3. Интерфейс с правильным компонентом и функцией ---
+# --- 4. Интерфейс ---
 iface = gr.Interface(
-    fn=gradio_inference, # <-- Используем нашу новую функцию-обертку
+    fn=gradio_inference,
     inputs=[
         gr.Sketchpad(
-            type="pil", # <-- Важно: получаем объект PIL.Image
+            type="pil", # Оставляем pil, чтобы Gradio присылал объекты Pillow
             label="Нарисуйте формулу"
         ),
         gr.Slider(
@@ -81,8 +94,9 @@ iface = gr.Interface(
     ],
     outputs=gr.Textbox(label="Предсказанная LaTeX-строка"),
     title="Handwritten Formula Recognition",
-    description="Нарисуйте формулу на холсте, чтобы получить ее в формате LaTeX. Модель лучше всего работает с одиночными, четко написанными символами и короткими формулами."
+    description="Нарисуйте формулу на холсте, чтобы получить ее в формате LaTeX."
 )
 
 if __name__ == "__main__":
+    # Используй share=True, если хочешь получить публичную ссылку
     iface.launch()