Update app.py

app.py

@@ -1,6 +1,12 @@
+# =====================================================
+# AI INSURANCE CLAIM GENERATOR (FINAL HIGH-ACCURACY VERSION)
+# Better TOTAL extraction + Better COMPANY extraction
+# Hugging Face Space Ready
+# =====================================================
+
 import gradio as gr
 import pytesseract
-from PIL import Image, ImageFilter, ImageOps
+from PIL import Image
 import torch
 import re
 import requests
@@ -12,7 +18,9 @@ from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 # CONFIG
 # =====================================================
 RESEND_API_KEY = os.getenv("RESEND_API_KEY")
-FROM_EMAIL = "AI Claims <claims@yudham.com>"   # verified sender
+
+# VERIFIED DOMAIN EMAIL
+FROM_EMAIL = "AI Claims <claims@yudham.com>"
 
 MODEL_NAME = "ngupta2026/sroie-layoutlm"
 
@@ -37,7 +45,7 @@ model.to(device)
 model.eval()
 
 # =====================================================
-# NORMALIZE BOUNDING BOXES
+# NORMALIZE BOX
 # =====================================================
 def normalize(box, width, height):
     return [
@@ -48,34 +56,72 @@ def normalize(box, width, height):
     ]
 
 # =====================================================
-# IMAGE PREPROCESSING (VERY IMPORTANT)
+# AVG CONF
+# =====================================================
+def avg_conf(lst):
+    if len(lst) == 0:
+        return 0
+    return sum(lst) / len(lst)
+
+# =====================================================
+# CLEAN MONEY
 # =====================================================
-def preprocess_image(image):
+def clean_amount(txt):
+    txt = txt.replace(",", "").replace("RM", "").replace("₹", "")
+    txt = txt.strip()
 
-    image = image.convert("RGB")
+    try:
+        val = float(txt)
+        return round(val, 2)
+    except:
+        return None
+
+# =====================================================
+# FIND BEST TOTAL (VERY IMPORTANT FIX)
+# =====================================================
+def extract_best_total(words):
+
+    candidates = []
+
+    for i, w in enumerate(words):
 
-    # upscale for OCR
-    w, h = image.size
-    image = image.resize((w * 2, h * 2))
+        # match amounts like:
+        # 102.40
+        # 1,234.50
+        # RM102.40
+        if re.fullmatch(r"(RM)?\d+[.,]?\d*\.\d{2}", w):
+            amt = clean_amount(w)
+            if amt:
+                candidates.append(amt)
 
-    # grayscale
-    image = image.convert("L")
+        elif re.fullmatch(r"\d+\.\d{2}", w):
+            amt = clean_amount(w)
+            if amt:
+                candidates.append(amt)
 
-    # sharpen
-    image = image.filter(ImageFilter.SHARPEN)
+    # choose sensible max under 100000
+    candidates = [x for x in candidates if 1 <= x <= 100000]
 
-    # auto contrast
-    image = ImageOps.autocontrast(image)
+    if len(candidates) == 0:
+        return "Not Found"
 
-    return image
+    return f"{max(candidates):.2f}"
 
 # =====================================================
-# CONFIDENCE AVG
+# COMPANY CLEANER
 # =====================================================
-def avg_conf(lst):
-    if len(lst) == 0:
-        return 0
-    return sum(lst) / len(lst)
+def clean_company(txt):
+
+    txt = txt.strip()
+
+    # remove garbage symbols
+    txt = re.sub(r"[^A-Za-z0-9&().,\- ]", "", txt)
+
+    # remove too short
+    if len(txt) < 3:
+        return "Not Found"
+
+    return txt
 
 # =====================================================
 # OCR + EXTRACTION
@@ -83,13 +129,12 @@ def avg_conf(lst):
 def extract_receipt(image):
 
     try:
-        image = preprocess_image(image)
+        image = image.convert("RGB")
+        image.thumbnail((1400, 1400))
 
-        # Better OCR mode for receipts
         data = pytesseract.image_to_data(
             image,
-            output_type=pytesseract.Output.DICT,
-            config="--oem 3 --psm 4"
+            output_type=pytesseract.Output.DICT
         )
 
         words = []
@@ -97,16 +142,16 @@ def extract_receipt(image):
 
         for i in range(len(data["text"])):
 
-            txt = data["text"][i].strip()
+            text = data["text"][i].strip()
 
-            if txt != "" and txt != "|":
+            if text != "" and len(text) > 1:
 
                 x = data["left"][i]
                 y = data["top"][i]
                 w = data["width"][i]
                 h = data["height"][i]
 
-                words.append(txt)
+                words.append(text)
                 boxes.append([x, y, x + w, y + h])
 
         if len(words) == 0:
@@ -115,9 +160,9 @@ def extract_receipt(image):
         width, height = image.size
         boxes = [normalize(b, width, height) for b in boxes]
 
-        # =================================================
+        # -------------------------------------------------
         # TOKENIZER
-        # =================================================
+        # -------------------------------------------------
         encoding = tokenizer(
             words,
             boxes=boxes,
@@ -130,9 +175,9 @@ def extract_receipt(image):
 
         encoding = {k: v.to(device) for k, v in encoding.items()}
 
-        # =================================================
-        # MODEL PREDICTION
-        # =================================================
+        # -------------------------------------------------
+        # MODEL
+        # -------------------------------------------------
         with torch.no_grad():
             outputs = model(**encoding)
 
@@ -141,92 +186,61 @@ def extract_receipt(image):
         preds = torch.argmax(probs, dim=2)[0][:len(words)]
         confs = torch.max(probs, dim=2)[0][0][:len(words)]
 
-        result = {
-            "company": [],
-            "date": [],
-            "total": []
-        }
+        company_words = []
+        company_conf = []
 
-        conf_store = {
-            "company": [],
-            "date": [],
-            "total": []
-        }
-
-        # =================================================
-        # EXTRACT ENTITIES
-        # =================================================
+        # -------------------------------------------------
+        # ENTITY EXTRACTION
+        # -------------------------------------------------
         for word, pred, conf in zip(words, preds, confs):
 
             label = id2label[pred.item()]
             c = conf.item()
 
-            # -------------------------
-            # COMPANY
-            # -------------------------
             if label == "COMPANY":
-                result["company"].append(word)
-                conf_store["company"].append(c)
-
-            # -------------------------
-            # DATE
-            # -------------------------
-            if re.search(r"\d{1,2}[/-]\d{1,2}[/-]\d{2,4}", word):
-                result["date"].append(word)
-                conf_store["date"].append(c)
-
-            # -------------------------
-            # TOTAL
-            # -------------------------
-            cleaned = word.replace(",", "").replace("₹", "").replace("$", "")
-
-            if re.fullmatch(r"\d+(\.\d{1,2})?", cleaned):
-
-                try:
-                    value = float(cleaned)
-
-                    # realistic receipt range
-                    if 1 <= value <= 10000:
-                        result["total"].append(value)
-                        conf_store["total"].append(c)
-
-                except:
-                    pass
-
-        # =================================================
-        # FINAL CLEANUP
-        # =================================================
+                company_words.append(word)
+                company_conf.append(c)
 
+        # -------------------------------------------------
         # COMPANY
-        company = " ".join(result["company"][:6]).strip()
-
-        if company == "":
-            # fallback top words
-            company = " ".join(words[:3])
+        # -------------------------------------------------
+        company = " ".join(company_words[:6]) if company_words else words[0]
+        company = clean_company(company)
 
+        # -------------------------------------------------
         # DATE
-        date = result["date"][0] if result["date"] else "Not Found"
+        # -------------------------------------------------
+        date = "Not Found"
+
+        for w in words:
+            if re.search(r"\d{2}[/-]\d{2}[/-]\d{2,4}", w):
+                date = w
+                break
 
-        # TOTAL = best realistic amount
-        if result["total"]:
-            total = f"{max(result['total']):.2f}"
-        else:
-            total = "Not Found"
+        # -------------------------------------------------
+        # TOTAL (NEW LOGIC)
+        # -------------------------------------------------
+        total = extract_best_total(words)
 
+        # -------------------------------------------------
         # CONFIDENCE
-        company_conf = avg_conf(conf_store["company"])
-        date_conf = avg_conf(conf_store["date"])
-        total_conf = avg_conf(conf_store["total"])
+        # -------------------------------------------------
+        conf = avg_conf(company_conf)
+
+        if total != "Not Found":
+            conf += 0.10
 
-        overall = (company_conf + date_conf + total_conf) / 3
+        conf = min(conf, 0.99)
 
-        return {
+        result = {
             "company": company,
             "date": date,
             "total": total,
-            "confidence": round(overall, 3)
+            "confidence": round(conf, 3)
         }
 
+        return result
+
     except Exception as e:
         return {"error": str(e)}
 
@@ -250,22 +264,18 @@ def decision_layer(conf):
 def send_claim_email(to_email, extracted):
 
     if not RESEND_API_KEY:
-        return "❌ Missing RESEND_API_KEY secret"
+        return "❌ Missing RESEND_API_KEY"
 
     subject = "Insurance Claim Request"
 
     html_body = f"""
     <h2>Insurance Claim Request</h2>
 
-    <p>Dear Claims Team,</p>
-
-    <p>Please process reimbursement request.</p>
+    <p><b>Provider:</b> {extracted['company']}</p>
+    <p><b>Date:</b> {extracted['date']}</p>
+    <p><b>Amount:</b> ₹{extracted['total']}</p>
 
-    <p><b>Provider Name:</b> {extracted['company']}</p>
-    <p><b>Bill Date:</b> {extracted['date']}</p>
-    <p><b>Claim Amount:</b> ₹{extracted['total']}</p>
-
-    <p>Regards,<br>AI Claims System</p>
+    <p>Regards,<br>AI Claims Bot</p>
     """
 
     try:
@@ -285,7 +295,7 @@ def send_claim_email(to_email, extracted):
         )
 
         if response.status_code in [200, 201]:
-            return f"✅ Email sent successfully to {to_email}"
+            return f"✅ Email sent to {to_email}"
 
         return f"❌ Email failed: {response.text}"
 
@@ -303,6 +313,7 @@ def process_and_send(image, email_id):
         return extracted, extracted["error"]
 
     conf = extracted["confidence"]
+
     decision = decision_layer(conf)
 
     extracted["decision"] = decision
@@ -311,10 +322,10 @@ def process_and_send(image, email_id):
         email_status = send_claim_email(email_id, extracted)
 
     elif decision == "REVIEW":
-        email_status = f"⚠️ Human review required (confidence={conf})"
+        email_status = f"⚠️ Manual review required ({conf})"
 
     else:
-        email_status = f"❌ Rejected (low confidence={conf})"
+        email_status = f"❌ Rejected ({conf})"
 
     return extracted, email_status
 
@@ -326,7 +337,7 @@ demo = gr.Interface(
 
     inputs=[
         gr.Image(type="pil", label="Upload Receipt"),
-        gr.Textbox(label="Enter Destination Email")
+        gr.Textbox(label="Destination Email")
     ],
 
     outputs=[
@@ -335,7 +346,7 @@ demo = gr.Interface(
     ],
 
     title="📄 AI Insurance Claim Generator",
-    description="Upload receipt → Extract fields accurately → Confidence Check → Auto Email"
+    description="Upload receipt → Better AI extraction → Confidence → Auto Email"
 )
 
 demo.launch()