Update ocr_engine.py

ocr_engine.py

@@ -10,7 +10,7 @@ model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwrit
 def clean_ocr_text(text):
     # Fix common OCR misreads
     text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0")
-    return re.sub(r"[^\d.kg]", "", text.lower())  # keep digits, dot, k, g
+    return re.sub(r"[^\d.kg]", "", text.lower())
 
 def restore_decimal(text):
     if re.fullmatch(r"\d{5}", text):
@@ -25,11 +25,11 @@ def extract_unit_from_text(raw_text):
         return "kg"
     elif "g" in raw_text:
         return "g"
-    return "g"  # fallback if unit not found
+    return "g"  # fallback if no unit
 
 def extract_weight(image):
     try:
-        # Enhance image
+        # Resize & sharpen image
         image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC)
         image = image.filter(ImageFilter.SHARPEN)
 
@@ -40,17 +40,24 @@ def extract_weight(image):
 
         cleaned = clean_ocr_text(raw_text)
 
-        # Try direct match (e.g., 52.25 kg or 250.5g)
+        # Case 1: Match decimal with unit
         match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned)
         if match:
             return f"{match.group(1)} {match.group(2) or ''}".strip(), raw_text
 
-        # Fallback if no decimal found: convert big number like 53255 to 52.255
-        fallback_match = re.search(r"\d{4,5}", cleaned)
-        if fallback_match:
-            decimal_fixed = restore_decimal(fallback_match.group())
-            return decimal_fixed, raw_text
+        # Case 2: Large number fallback like 53255 → 52.255
+        match = re.search(r"\d{4,5}", cleaned)
+        if match:
+            decimal_fixed = restore_decimal(match.group())
+            unit = extract_unit_from_text(raw_text)
+            return f"{decimal_fixed} {unit}", raw_text
+
+        # Final fallback: plain number
+        match = re.search(r"\d+", cleaned)
+        if match:
+            return f"{match.group()} g", raw_text
 
         return "Error: No valid weight found", raw_text
+
     except Exception as e:
         return f"Error: {str(e)}", ""