Update ocr_engine.py

ocr_engine.py

@@ -3,21 +3,31 @@ from PIL import Image
 import torch
 import re
 
-# Load TrOCR model and processor
+# Load TrOCR processor and model once
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
+def clean_ocr_text(text):
+    # Fix common OCR mistakes
+    text = text.replace(",", ".")               # comma to dot
+    text = re.sub(r"[^\d\.kg]", "", text.lower())  # keep only digits, dot, k, g
+    return text
+
 def extract_weight(image):
     try:
-        # Step 1: OCR with TrOCR
+        # TrOCR inference
         pixel_values = processor(images=image, return_tensors="pt").pixel_values
         generated_ids = model.generate(pixel_values)
-        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
-        print("OCR Output:", text)
+        raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        print("OCR Raw Output:", raw_text)
+
+        # Clean and normalize text
+        cleaned_text = clean_ocr_text(raw_text)
+        print("Cleaned OCR:", cleaned_text)
 
-        # Step 2: Search for weight pattern (g or kg)
-        pattern = r'(\d{1,5}(?:\.\d{1,3})?)\s*(kg|g)'  # e.g., 75.25 g, 98.78kg
-        match = re.search(pattern, text.lower())
+        # Flexible regex to catch even minor issues (e.g., 52.2g, 98.7kg)
+        pattern = r'(\d{1,5}(?:\.\d{1,3})?)\s*(kg|g)'
+        match = re.search(pattern, cleaned_text)
 
         if match:
             value = match.group(1)