Update utils/image_processor.py

utils/image_processor.py

@@ -1,58 +1,61 @@
-import torch
+import os
 import cv2
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import numpy as np
 from PIL import Image
-import os
+import torch
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 
-# Load processor and model only once
+# ✅ Load TrOCR model and processor once
 try:
     processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
     model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
-    model.to("cpu")  # You can set to "cuda" if running with GPU
-    print("✅ TrOCR model loaded successfully.")
-    ocr_available = True
+    model.eval()
+    print("✅ TrOCR model loaded.")
+    trocr_available = True
 except Exception as e:
-    print(f"❌ TrOCR initialization failed: {str(e)}")
-    ocr_available = False
+    print(f"❌ Failed to load TrOCR: {e}")
+    trocr_available = False
+
+def preprocess_image(image):
+    """
+    Preprocess image for OCR: convert to grayscale and enhance contrast.
+    """
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    denoised = cv2.fastNlMeansDenoising(gray, h=10)
+    processed = cv2.adaptiveThreshold(
+        denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+        cv2.THRESH_BINARY, 11, 2
+    )
+    return processed
 
 def extract_text_from_image(image_path):
     """
-    Extract text from an image file using EasyOCR or TrOCR (if integrated).
+    Extract text using TrOCR from a handwritten image
     """
     try:
+        if not trocr_available:
+            raise ValueError("TrOCR is not available.")
+
         print(f"📂 Reading image from: {image_path}")
 
-        if not ocr_available:
-            raise ValueError("EasyOCR is not available")
-        
+        # Read and preprocess image
         image = cv2.imread(image_path)
         if image is None:
-            raise ValueError(f"Could not read image at {image_path}")
-        
-        processed_image = preprocess_image(image)
-
-        temp_path = os.path.join(os.path.dirname(image_path), f"temp_{os.path.basename(image_path)}")
-        cv2.imwrite(temp_path, processed_image)
+            raise ValueError(f"Could not load image: {image_path}")
 
-        results = reader.readtext(temp_path)
-
-        try:
-            os.remove(temp_path)
-        except:
-            pass
+        processed_image = preprocess_image(image)
+        pil_image = Image.fromarray(processed_image).convert("RGB")
 
-        text = ' '.join([result[1] for result in results]).strip()
+        # TrOCR expects pixel values between 0-1
+        pixel_values = processor(images=pil_image, return_tensors="pt").pixel_values
 
-        if not text:
-            results = reader.readtext(image_path)
-            text = ' '.join([result[1] for result in results]).strip()
+        # Disable gradient for inference
+        with torch.no_grad():
+            generated_ids = model.generate(pixel_values)
+            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # ✅ Log extracted text
-        print(f"📝 Extracted text from {os.path.basename(image_path)}:\n{text}\n")
+        return generated_text.strip() or "Text extraction failed. Please enter text manually."
 
-        return text or "Text extraction failed. Please enter text manually."
-    
     except Exception as e:
-        print(f"OCR failed: {str(e)}")
+        print(f"OCR failed: {e}")
         return "Text extraction failed. Please enter text manually."
-