Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from PIL import Image
 from datetime import datetime
 import pytz
-from ocr_engine import extract_weight
+from ocr_engine import extract_weight, extract_unit_from_text
 from simple_salesforce import Salesforce
 import base64
 import re
@@ -27,26 +27,28 @@ def process_image(image):
     if image is None:
         return "❌ No image provided", "", None, gr.update(visible=True)
     try:
-        weight = extract_weight(image)
+        result = extract_weight(image)
+        weight, raw_text = result if isinstance(result, tuple) else (result, "")
+
         print("🧠 Final OCR Result:", weight)
+        print("🔤 OCR Raw Text:", raw_text)
 
         ist = pytz.timezone('Asia/Kolkata')
         timestamp = datetime.now(ist).strftime("%Y-%m-%d %H:%M:%S IST")
 
-        if not weight or weight.startswith("Error"):
+        if not weight or (isinstance(weight, str) and weight.startswith("Error")):
             return f"❌ OCR Error: {weight}", "", image, gr.update(visible=True)
 
         match = re.search(r'(\d{1,3}\.\d{1,3})\s*(kg|g)?', weight)
         if match:
             numeric_value = float(match.group(1))
-            unit = match.group(2) if match.group(2) else "g"
+            unit = match.group(2) if match.group(2) else extract_unit_from_text(raw_text)
         else:
-            # fallback try to fix raw number like 53255 → 52.255
             cleaned = re.sub(r"[^\d]", "", weight)
             decimal_fixed = restore_decimal(cleaned)
             try:
                 numeric_value = float(decimal_fixed)
-                unit = "g"
+                unit = extract_unit_from_text(raw_text)
             except:
                 return f"❌ Could not extract number | OCR: {weight}", "", image, gr.update(visible=True)
 
@@ -85,7 +87,7 @@ def process_image(image):
     except Exception as e:
         return f"Error: {str(e)}", "", None, gr.update(visible=True)
 
-with gr.Blocks(css=".gr-button {background-color: #2e7d32 !important; color: white !important;}") as demo:
+with gr.Blocks(css=".gr-button {background-color: #2e7d32 !important; color: white !important;}" ) as demo:
     gr.Markdown("""
     <h1 style='text-align: center; color: #2e7d32;'>📷 Auto Weight Logger</h1>
     <p style='text-align: center;'>Upload or capture a digital weight image. Detects weight using AI OCR and logs it into Salesforce.</p>