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("", self.start_draw) self.canvas.bind("", 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()