main.py

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()