Update ocr_engine.py

ocr_engine.py

@@ -3,14 +3,14 @@ from PIL import Image
 import torch
 import re
 
-# Load TrOCR model and processor once
+# Load model and processor
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
 def clean_ocr_text(text):
     print("[RAW OCR]", text)
     text = text.replace(",", ".").replace("s", "5").replace("o", "0").lower()
-    text = re.sub(r"[^\d\.kg]", "", text)  # Keep digits, dot, kg
+    text = re.sub(r"[^\d\.kg]", "", text)
     print("[CLEANED OCR]", text)
     return text
 
@@ -22,11 +22,16 @@ def extract_weight(image):
 
         cleaned = clean_ocr_text(raw_text)
 
-        # Regex for weight
+        # First try with unit
         match = re.search(r'(\d{1,5}(?:\.\d{1,3})?)\s*(kg|g)', cleaned)
         if match:
             return f"{match.group(1)} {match.group(2)}"
-        else:
-            return f"No valid weight found | OCR: {cleaned}"
+
+        # Fallback: only number, assume grams
+        fallback = re.search(r'(\d{1,5}(?:\.\d{1,3})?)', cleaned)
+        if fallback:
+            return f"{fallback.group(1)} g"
+
+        return f"No valid weight found | OCR: {cleaned}"
     except Exception as e:
         return f"Error: {str(e)}"