Update ocr_engine.py

ocr_engine.py

@@ -3,16 +3,20 @@ from PIL import Image, ImageFilter
 import torch
 import re
 
-# Load model
 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").replace("O", "0")
-    text = re.sub(r"[^\d.kg]", "", text.lower())  # Keep digits, dots, and kg
-    print("[CLEANED OCR]", text)
-    return text
+    return re.sub(r"[^\d.kg]", "", text.lower())
+
+def extract_unit_from_text(raw_text):
+    raw = raw_text.lower()
+    if "kg" in raw:
+        return "kg"
+    elif "g" in raw:
+        return "g"
+    return "g"  # default fallback
 
 def restore_decimal(text):
     if re.fullmatch(r"\d{5}", text):
@@ -32,20 +36,15 @@ def extract_weight(image):
 
         cleaned = clean_ocr_text(raw_text)
 
-        # Try direct match: e.g., 52.25 kg or 75.0 g
-        match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)", cleaned)
+        match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned)
         if match:
-            return f"{match.group(1)} {match.group(2)}"
-
-        # Try fallback: extract digits and manually guess decimal
-        fallback_match = re.search(r"(\d{4,5})", cleaned)
-        if fallback_match:
-            fallback_value = restore_decimal(fallback_match.group(1))
+            return f"{match.group(1)} {match.group(2) or ''}", raw_text
 
-            # Check for presence of unit hints in raw_text
-            unit = "kg" if "kg" in raw_text.lower() else "g"
-            return f"{fallback_value} {unit}"
+        fallback = re.search(r"\d{4,5}", cleaned)
+        if fallback:
+            fixed = restore_decimal(fallback.group())
+            return f"{fixed}", raw_text
 
-        return f"No valid weight found | OCR: {cleaned}"
+        return f"No valid weight found | OCR: {cleaned}", raw_text
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}", ""