ocrfood / app /utils.py
ZienabM's picture
Update app/utils.py
b052766 verified
# utils.py
import os
from transformers import AutoModel, AutoTokenizer
from PIL import Image, ImageEnhance, ImageFilter
import torch
import logging
logger = logging.getLogger(__name__)
class OCRModel:
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super(OCRModel, cls).__new__(cls)
cls._instance.initialize()
return cls._instance
def initialize(self):
try:
logger.info("Initializing OCR model...")
# تهيئة النموذج والتوكينايزر
self.tokenizer = AutoTokenizer.from_pretrained('RufusRubin777/GOT-OCR2_0_CPU', trust_remote_code=True)
self.model = AutoModel.from_pretrained(
'RufusRubin777/GOT-OCR2_0_CPU',
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 = "cpu" # هذا النموذج مصمم للـ CPU
self.model = self.model.eval().cpu()
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:
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)
# استخدام النموذج للتعرف على النص
try:
result = self.model.chat(self.tokenizer, temp_image_path, ocr_type='format')
logger.info(f"Successfully extracted text: {result[:100]}...")
# حذف الملف المؤقت
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)
return f"Error processing image: {str(e)}"