Update app.py

app.py

@@ -1,90 +1,58 @@
 import os
-import gc
-import psutil
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-
-class MultiModelSystem:
-    """
-    سیستم چندمدلی با مدیریت حافظه و ذخیره‌سازی موقت مدل‌ها در دیسک.
-    """
-
-    def __init__(self, memory_limit_gb=15):
-        self.models = {}
-        self.memory_limit_gb = memory_limit_gb
-        self.model_cache_dir = "model_cache"
-        os.makedirs(self.model_cache_dir, exist_ok=True)
-
-    def check_memory_usage(self):
-        mem = psutil.virtual_memory()
-        used_gb = mem.used / (1024 ** 3)
-        print(f"Memory usage: {mem.percent}% ({used_gb:.2f} GB used)")
-        if used_gb > self.memory_limit_gb:
-            raise MemoryError(f"Memory limit exceeded: {used_gb:.2f} GB used (limit: {self.memory_limit_gb} GB)")
-
-    def load_model(self, task, model_id):
-        """
-        بارگذاری مدل از کش یا ذخیره‌سازی.
-        """
-        cache_path = os.path.join(self.model_cache_dir, f"{task}.bin")
-        if task not in self.models:
-            self.check_memory_usage()
-            print(f"Loading model for task '{task}'...")
-            if os.path.exists(cache_path):
-                print(f"Loading model from cache: {cache_path}")
-                self.models[task] = joblib.load(cache_path)
-            else:
-                model = AutoModelForSeq2SeqLM.from_pretrained(
-                    model_id, torch_dtype="auto", low_cpu_mem_usage=True
-                )
-                tokenizer = AutoTokenizer.from_pretrained(model_id)
-                self.models[task] = pipeline("translation", model=model, tokenizer=tokenizer)
-                joblib.dump(self.models[task], cache_path)
-                print(f"Model cached at {cache_path}")
-
-    def unload_model(self, task):
-        """
-        تخلیه مدل از حافظه.
-        """
-        if task in self.models:
-            print(f"Unloading model for task '{task}'...")
-            del self.models[task]
-            gc.collect()
-
-    def process_task(self, task, model_id, **kwargs):
-        """
-        پردازش وظیفه با بارگذاری موقت مدل.
-        """
-        self.load_model(task, model_id)
-        model = self.models[task]
-        if task == "translation":
-            text = kwargs.get("text", "")
-            return model(text)
-        elif task == "qa":
-            question = kwargs.get("question", "")
-            context = kwargs.get("context", "")
-            return model(question=question, context=context)
-        else:
-            raise ValueError(f"Unsupported task: {task}")
+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__":
-    MODEL_CONFIG = {
-        "translation": "Helsinki-NLP/opus-mt-en-ro",  # مدل سبک‌تر
-        "qa": "distilbert-base-uncased-distilled-squad",  # مدل فشرده
-    }
-
-    tasks = [
-        {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!"}},
-        {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "What is AI?", "context": "AI is artificial intelligence."}}
-    ]
-
-    system = MultiModelSystem(memory_limit_gb=15)
-
-    for task_info in tasks:
-        try:
-            system.check_memory_usage()
-            result = system.process_task(task_info["task"], task_info["model_id"], **task_info["kwargs"])
-            print(f"Result for task '{task_info['task']}':", result)
-        except Exception as e:
-            print(f"Error during task '{task_info['task']}':", str(e))
-        finally:
-            system.unload_model(task_info["task"])
+    app = OCRApp()
+    app.mainloop()