Create main.py

main.py

@@ -0,0 +1,179 @@
+import tkinter as tk
+from tkinter import messagebox
+from PIL import Image, ImageDraw, ImageOps
+import os
+import torch
+import pickle
+import numpy as np
+import torchvision.transforms as transforms
+from model import PosFormer, Vocab, PosVocab, ScaleToLimitRange
+
+class FormulaApp:
+    def __init__(self, root):
+        self.root = root
+        self.root.title("Распознавание математических формул")
+
+        # Увеличим размер холста для лучшего распознавания
+        self.canvas_width = 800
+        self.canvas_height = 400
+        self.pen_color = "black"
+
+        # Загрузка модели и словаря
+        self.model = None
+        self.main_vocab = None
+        self.load_model()
+
+        # Создание интерфейса
+        self.canvas = tk.Canvas(root, width=self.canvas_width, height=self.canvas_height, bg='white')
+        self.canvas.pack(pady=10)
+
+        self.image = Image.new("RGB", (self.canvas_width, self.canvas_height), "white")
+        self.draw = ImageDraw.Draw(self.image)
+
+        self.canvas.bind("<Button-1>", self.start_draw)
+        self.canvas.bind("<B1-Motion>", self.draw_on_canvas)
+
+        btn_frame = tk.Frame(root)
+        btn_frame.pack(pady=5)
+
+        tk.Button(btn_frame, text="Распознать", command=self.recognize).pack(side=tk.LEFT, padx=5)
+        tk.Button(btn_frame, text="Очистить", command=self.clear_canvas).pack(side=tk.LEFT, padx=5)
+        tk.Button(btn_frame, text="Сохранить", command=self.save_all).pack(side=tk.LEFT, padx=5)
+
+        self.eraser_button = tk.Button(btn_frame, text="Ластик", command=self.toggle_eraser)
+        self.eraser_button.pack(side=tk.LEFT, padx=5)
+
+        self.result_label = tk.Label(root, text="Результат: ", font=("Arial", 12))
+        self.result_label.pack(pady=10)
+
+        self.last_x = None
+        self.last_y = None
+        self.is_erasing = False
+
+        # Папки для хранения
+        self.draw_folder = "рисунки"
+        self.text_folder = "формулы"
+        os.makedirs(self.draw_folder, exist_ok=True)
+        os.makedirs(self.text_folder, exist_ok=True)
+
+        # Инициализация препроцессинга
+        self.size_controller = ScaleToLimitRange(h_lo=32, h_hi=256, w_lo=32, w_hi=512)
+        self.to_tensor = transforms.ToTensor()
+        self.normalize = transforms.Normalize([0.5], [0.5])
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def load_model(self):
+        try:
+            print("Загрузка словаря...")
+            with open("crohme_vocab.pkl", "rb") as f:
+                self.main_vocab = pickle.load(f)
+            pos_vocab = PosVocab()
+            print("Словарь загружен.")
+
+            print("Создание модели...")
+            self.model = PosFormer(self.main_vocab, pos_vocab).to(self.device)
+            print("Модель создана.")
+
+            print("Загрузка весов из 43%.pth...")
+            checkpoint = torch.load("43%.pth", map_location=self.device)
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+            self.model.eval()
+            print(" Модель успешно загружена и готова к работе.")
+
+        except Exception as e:
+            print(f" ОШИБКА: Не удалось загрузить модель или словарь: {e}")
+            messagebox.showerror("Ошибка", f"Не удалось загрузить модель: {e}")
+
+    def preprocess_image(self, pil_image):
+        """Подготовка изображения для модели"""
+        # Конвертируем в grayscale и инвертируем цвета
+        img_gray = pil_image.convert('L')
+        img_gray = ImageOps.invert(img_gray)
+        
+        # Конвертируем в numpy для обработки
+        img_np = np.array(img_gray)
+        
+        # Масштабируем изображение
+        img_np = self.size_controller(img_np)
+        
+        # Конвертируем обратно в PIL и применяем трансформации
+        final_img_pil = Image.fromarray(img_np)
+        tensor = self.normalize(self.to_tensor(final_img_pil))
+        
+        # Добавляем batch dimension
+        return tensor.unsqueeze(0).to(self.device)
+
+    def recognize(self):
+        if self.model is None:
+            messagebox.showerror("Ошибка", "Модель не загружена")
+            return
+
+        try:
+            # Препроцессинг изображения
+            image_tensor = self.preprocess_image(self.image)
+            
+            # Распознавание с beam search
+            predicted_ids = self.model.generate_beam_search(image_tensor, beam_size=10, max_gen_len=200)
+            
+            # Декодирование результата
+            latex_string = self.main_vocab.decode(predicted_ids[0].cpu().tolist())
+            self.result_label.config(text=f"Результат: $${latex_string}$$")
+            
+        except Exception as e:
+            messagebox.showerror("Ошибка", f"Ошибка при распознавании: {e}")
+            self.result_label.config(text="Результат: ошибка распознавания")
+
+    def start_draw(self, event):
+        self.last_x = event.x
+        self.last_y = event.y
+
+    def draw_on_canvas(self, event):
+        x, y = event.x, event.y
+        color = "white" if self.is_erasing else self.pen_color
+        width = 8 if self.is_erasing else 4
+
+        self.canvas.create_line(self.last_x, self.last_y, x, y, width=width,
+                              fill=color, capstyle=tk.ROUND, smooth=True)
+        self.draw.line([self.last_x, self.last_y, x, y], fill=color, width=width)
+
+        self.last_x = x
+        self.last_y = y
+
+    def toggle_eraser(self):
+        self.is_erasing = not self.is_erasing
+        self.eraser_button.config(text="Карандаш" if self.is_erasing else "Ластик")
+
+    def clear_canvas(self):
+        self.canvas.delete("all")
+        self.image = Image.new("RGB", (self.canvas_width, self.canvas_height), "white")
+        self.draw = ImageDraw.Draw(self.image)
+        self.result_label.config(text="Результат: ")
+        self.is_erasing = False
+        self.eraser_button.config(text="Ластик")
+
+    def get_next_filename(self, folder, prefix, ext):
+        i = 1
+        while True:
+            filename = f"{prefix}_{i}.{ext}"
+            full_path = os.path.join(folder, filename)
+            if not os.path.exists(full_path):
+                return full_path
+            i += 1
+
+    def save_all(self):
+        # Сохранить рисунок
+        image_path = self.get_next_filename(self.draw_folder, "рисунок", "png")
+        self.image.save(image_path)
+
+        # Сохранить распознанную формулу
+        text = self.result_label.cget("text").replace("Результат: ", "")
+        formula_path = self.get_next_filename(self.text_folder, "формула", "txt")
+        with open(formula_path, "w", encoding="utf-8") as f:
+            f.write(text)
+
+        messagebox.showinfo("Сохранение", f"Изображение и формула сохранены:\n{image_path}\n{formula_path}")
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    app = FormulaApp(root)
+    root.mainloop()