import os
from datasets import load_dataset
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import customtkinter as ctk
from tkinter import filedialog
from PIL import Image, ImageTk


dataset = load_dataset("hezarai/parsynth-ocr-200k")
# تنظیم GUI
ctk.set_appearance_mode("System")
ctk.set_default_color_theme("blue")

class OCRApp(ctk.CTk):
    def __init__(self):
        super().__init__()
        self.title("OCR with Hugging Face")
        self.geometry("800x600")

        # مدل و پردازشگر
        self.processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
        self.model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

        # عناصر رابط کاربری
        self.image_label = ctk.CTkLabel(self, text="No Image Selected", width=400, height=300, corner_radius=8)
        self.image_label.pack(pady=20)

        self.upload_button = ctk.CTkButton(self, text="Upload Image", command=self.upload_image)
        self.upload_button.pack(pady=10)

        self.result_label = ctk.CTkTextbox(self, height=200)
        self.result_label.pack(pady=10, fill="both", expand=True)

    def upload_image(self):
        file_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png;*.jpg;*.jpeg")])
        if not file_path:
            return

        # نمایش تصویر
        image = Image.open(file_path).resize((400, 300))
        self.image_tk = ImageTk.PhotoImage(image)
        self.image_label.configure(image=self.image_tk, text="")

        # انجام OCR
        text = self.perform_ocr(file_path)
        self.result_label.delete("1.0", "end")
        self.result_label.insert("1.0", text)

    def perform_ocr(self, image_path):
        image = Image.open(image_path).convert("RGB")
        pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
        generated_ids = self.model.generate(pixel_values)
        generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        return generated_text

if __name__ == "__main__":
    app = OCRApp()
    app.mainloop()