Update app/utils.py

app/utils.py

@@ -1,6 +1,6 @@
 # utils.py
 import os
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import AutoModel, AutoTokenizer
 from PIL import Image, ImageEnhance, ImageFilter
 import torch
 import logging
@@ -20,16 +20,20 @@ class OCRModel:
         try:
             logger.info("Initializing OCR model...")
             
-            # تهيئة النموذج والمعالج
-            self.processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
-            self.model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
+            # تهيئة النموذج والتوكينايزر
+            self.tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
+            self.model = AutoModel.from_pretrained(
+                'ucaslcl/GOT-OCR2_0',
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                device_map='cpu',
+                use_safetensors=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
             
-            # تحديد الجهاز
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
-            logger.info(f"Using device: {self.device}")
-            
-            self.model.to(self.device)
-            self.model.eval()
+            # تحديد الجهاز وتهيئة النموذج للتقييم
+            self.device = "cpu"  # هذا النموذج مصمم للـ CPU
+            self.model = self.model.eval().cpu()
             
             logger.info("Model initialization completed successfully")
             
@@ -68,33 +72,33 @@ class OCRModel:
         try:
             logger.info("Starting image processing")
             
+            # حفظ الصورة مؤقتاً لأن النموذج يتطلب مسار ملف
+            temp_image_path = "temp_image.jpg"
+            
             # إعادة تعيين مؤشر البداية للـ BytesIO
             image_stream.seek(0)
             
-            # فتح الصورة
+            # فتح وحفظ الصورة مؤقتاً
             image = Image.open(image_stream).convert('RGB')
-            
-            # تطبيق المعالجة المسبقة
             processed_image = self.preprocess_image(image)
+            processed_image.save(temp_image_path)
             
-            # معالجة الصورة للنموذج
-            pixel_values = self.processor(processed_image, return_tensors="pt").pixel_values.to(self.device)
-            
-            # التعرف على النص
-            with torch.no_grad():
-                generated_ids = self.model.generate(
-                    pixel_values,
-                    max_length=128,
-                    num_beams=4,
-                    length_penalty=2.0,
-                    early_stopping=True
-                )
+            # استخدام النموذج للتعرف على النص
+            try:
+                result = self.model.chat(self.tokenizer, temp_image_path, ocr_type='format')
+                logger.info(f"Successfully extracted text: {result[:100]}...")
                 
-            # تحويل النتيجة إلى نص
-            generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            
-            logger.info(f"Successfully extracted text: {generated_text[:100]}...")
-            return generated_text.strip()
+                # حذف الملف المؤقت
+                if os.path.exists(temp_image_path):
+                    os.remove(temp_image_path)
+                
+                return result.strip()
+                
+            except Exception as e:
+                logger.error(f"Error in OCR processing: {str(e)}", exc_info=True)
+                if os.path.exists(temp_image_path):
+                    os.remove(temp_image_path)
+                raise
             
         except Exception as e:
             logger.error(f"Error in image processing: {str(e)}", exc_info=True)