Update app/utils.py

app/utils.py

@@ -1,8 +1,12 @@
+# utils.py
 import os
-from transformers import AutoModel, AutoTokenizer
-from PIL import Image
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from PIL import Image, ImageEnhance, ImageFilter
 import torch
- 
+import logging
+
+logger = logging.getLogger(__name__)
+
 class OCRModel:
     _instance = None
 
@@ -13,34 +17,85 @@ class OCRModel:
         return cls._instance
 
     def initialize(self):
-        # تحميل النموذج مرة واحدة وتخزينه محلياً
-        model_path = os.getenv('MODEL_PATH', 'RufusRubin777/GOT-OCR2_0_CPU')
-        
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            model_path, 
-            trust_remote_code=True,
-            local_files_only=False  # سيتم تحميل الملفات إذا لم تكن موجودة
-        )
-        
-        self.model = AutoModel.from_pretrained(
-            model_path,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-            device_map='cpu',  # سيختار أفضل جهاز متاح
-            use_safetensors=True,
-            pad_token_id=self.tokenizer.eos_token_id
-        )
-        
-        self.model = self.model.eval()
-
-        
+        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.device = "cuda" if torch.cuda.is_available() else "cpu"
+            logger.info(f"Using device: {self.device}")
+            
+            self.model.to(self.device)
+            self.model.eval()
+            
+            logger.info("Model initialization completed successfully")
+            
+        except Exception as e:
+            logger.error(f"Error initializing model: {str(e)}", exc_info=True)
+            raise
+
+    def preprocess_image(self, image):
+        """معالجة مسبقة للصورة لتحسين جودة التعرف على النص"""
+        try:
+            # تحويل الصورة إلى RGB إذا لم تكن كذلك
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+
+            # تحسين التباين
+            enhancer = ImageEnhance.Contrast(image)
+            image = enhancer.enhance(1.5)
+
+            # تحسين الحدة
+            enhancer = ImageEnhance.Sharpness(image)
+            image = enhancer.enhance(1.5)
+
+            # تحسين السطوع
+            enhancer = ImageEnhance.Brightness(image)
+            image = enhancer.enhance(1.2)
+
+            # تطبيق فلتر لتنعيم الصورة قليلاً
+            image = image.filter(ImageFilter.SMOOTH)
+
+            return image
+        except Exception as e:
+            logger.error(f"Error in image preprocessing: {str(e)}", exc_info=True)
+            raise
+
     def process_image(self, image_stream):
         try:
-            # فتح الصورة من الذاكرة
-            image = Image.open(image_stream)
+            logger.info("Starting image processing")
             
+            # إعادة تعيين مؤشر البداية للـ BytesIO
+            image_stream.seek(0)
+            
+            # فتح الصورة
+            image = Image.open(image_stream).convert('RGB')
+            
+            # تطبيق المعالجة المسبقة
+            processed_image = self.preprocess_image(image)
+            
+            # معالجة الصورة للنموذج
+            pixel_values = self.processor(processed_image, return_tensors="pt").pixel_values.to(self.device)
+            
+            # التعرف على النص
             with torch.no_grad():
-                result = self.model.chat(self.tokenizer, image, ocr_type='format')
-            return result
+                generated_ids = self.model.generate(
+                    pixel_values,
+                    max_length=128,
+                    num_beams=4,
+                    length_penalty=2.0,
+                    early_stopping=True
+                )
+                
+            # تحويل النتيجة إلى نص
+            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()
+            
         except Exception as e:
+            logger.error(f"Error in image processing: {str(e)}", exc_info=True)
             return f"Error processing image: {str(e)}"