Update app.py

app.py

@@ -1,187 +1,262 @@
+import gradio as gr
+import pytesseract
 from PIL import Image
 import torch
 import re
+import requests
 import os
-import smtplib
-
-from email.mime.text import MIMEText
-from email.mime.multipart import MIMEMultipart
 
 from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 
 # =====================================================
-# LABELS
+# CONFIG
 # =====================================================
-label2id = {"O": 0, "COMPANY": 1, "DATE": 2, "TOTAL": 3}
+RESEND_API_KEY = os.getenv("RESEND_API_KEY")
+
+# Use verified sender from Resend
+FROM_EMAIL = "AI Claims <claims@yudham.com>"
+
+MODEL_NAME = "ngupta2026/sroie-layoutlm"
+
+label2id = {
+    "O": 0,
+    "COMPANY": 1,
+    "DATE": 2,
+    "TOTAL": 3
+}
+
 id2label = {v: k for k, v in label2id.items()}
 
 # =====================================================
 # LOAD MODEL
 # =====================================================
-MODEL_NAME = "ngupta2026/sroie-layoutlm"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 model = LayoutLMForTokenClassification.from_pretrained(MODEL_NAME)
 tokenizer = LayoutLMTokenizerFast.from_pretrained(MODEL_NAME)
+
 model.to(device)
 model.eval()
 
 # =====================================================
-# EMAIL CONFIG
-# Add these in Hugging Face Space Secrets:
-# EMAIL_USER = yourgmail@gmail.com
-# EMAIL_PASS = your_app_password
-# =====================================================
-EMAIL_USER = os.getenv("EMAIL_USER")
-EMAIL_PASS = os.getenv("EMAIL_PASS")
-
-# =====================================================
-# NORMALIZE BOXES
+# NORMALIZE BOX
 # =====================================================
 def normalize(box, width, height):
     return [
         int(1000 * box[0] / width),
+        int(1000 * box[1] / height),
+        int(1000 * box[2] / width),
         int(1000 * box[3] / height),
     ]
 
 # =====================================================
-# EXTRACT DATA
+# AVG CONFIDENCE
+# =====================================================
+def avg_conf(values):
+    if len(values) == 0:
+        return 0
+    return sum(values) / len(values)
+
+# =====================================================
+# OCR + EXTRACTION (IMPROVED ACCURACY)
 # =====================================================
 def extract_receipt(image):
 
-    data = pytesseract.image_to_data(
-        image,
-        output_type=pytesseract.Output.DICT
-    )
+    try:
+        # Keep quality high for OCR
+        image = image.convert("RGB")
+
+        data = pytesseract.image_to_data(
+            image,
+            output_type=pytesseract.Output.DICT,
+            config="--oem 3 --psm 6"
+        )
+
+        words = []
+        boxes = []
+
+        for i in range(len(data["text"])):
+
+            text = data["text"][i].strip()
+
+            if text != "" and text != "|":
 
-    words = []
-    boxes = []
+                x = data["left"][i]
+                y = data["top"][i]
+                w = data["width"][i]
+                h = data["height"][i]
 
-    for i in range(len(data["text"])):
+                words.append(text)
+                boxes.append([x, y, x + w, y + h])
 
-        text = data["text"][i].strip()
+        if len(words) == 0:
+            return {"error": "No text detected"}
 
-        if text != "":
-            h = data["height"][i]
+        width, height = image.size
+        boxes = [normalize(b, width, height) for b in boxes]
 
-            words.append(text)
-            boxes.append([x, y, x + w, y + h])
+        # IMPORTANT: use 512 for better predictions
+        encoding = tokenizer(
+            words,
+            boxes=boxes,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            is_split_into_words=True,
+            max_length=512
+        )
+
+        encoding = {k: v.to(device) for k, v in encoding.items()}
+
+        with torch.no_grad():
+            outputs = model(**encoding)
 
-    if len(words) == 0:
-        return {"error": "No text detected"}
+        probs = torch.softmax(outputs.logits, dim=2)
 
+        preds = torch.argmax(probs, dim=2)[0][:len(words)]
+        confs = torch.max(probs, dim=2)[0][0][:len(words)]
 
-    width, height = image.size
-    boxes = [normalize(box, width, height) for box in boxes]
+        result = {
+            "company": [],
+            "date": [],
+            "total": []
+        }
 
+        conf_store = {
+            "company": [],
+            "date": [],
+            "total": []
+        }
 
-    encoding = tokenizer(
-        words,
-        boxes=boxes,
-        max_length=512
-    )
+        # =================================================
+        # TOKEN LEVEL EXTRACTION
+        # =================================================
+        for word, pred, conf in zip(words, preds, confs):
 
-    encoding = {k: v.to(device) for k, v in encoding.items()}
+            label = id2label[pred.item()]
+            c = conf.item()
 
+            # COMPANY from model
+            if label == "COMPANY":
+                result["company"].append(word)
+                conf_store["company"].append(c)
 
-    with torch.no_grad():
-        outputs = model(**encoding)
+            # DATE regex
+            if re.search(r"\d{1,2}[/-]\d{1,2}[/-]\d{2,4}", word):
+                result["date"].append(word)
+                conf_store["date"].append(c)
 
-    predictions = torch.argmax(outputs.logits, dim=2)[0][:len(words)]
+            # TOTAL numeric values
+            cleaned = word.replace(",", "").replace("₹", "")
 
-    result = {
-        "company": [],
-        "date": [],
-        "total": []
-    }
+            if re.fullmatch(r"\d+(\.\d{1,2})?", cleaned):
+                try:
+                    value = float(cleaned)
 
-    for word, pred in zip(words, predictions):
+                    # Better range for totals
+                    if value >= 10:
+                        result["total"].append(value)
+                        conf_store["total"].append(c)
 
-        label = id2label[pred.item()]
+                except:
+                    pass
 
-        # company from model
-        if label == "COMPANY":
-            result["company"].append(word)
+        # =================================================
+        # FINAL CLEANUP
+        # =================================================
 
-        # date from regex
-        if re.search(r"\d{2}[/-]\d{2}[/-]\d{2,4}", word):
-            result["date"].append(word)
+        # COMPANY
+        company = " ".join(result["company"][:6]).strip()
+        if company == "":
+            company = "Not Found"
 
-        # total from regex
-        if re.search(r"\d+(\.\d{2})?", word):
-            try:
-                value = float(word.replace(",", ""))
-                if value > 50:
-                    result["total"].append(word)
-            except:
-                pass
+        # DATE
+        date = result["date"][0] if result["date"] else "Not Found"
 
-    result["company"] = (
-        " ".join(result["company"])
-        if result["company"] else "Not Found"
-    )
+        # TOTAL = highest amount (better than last token)
+        total = str(max(result["total"])) if result["total"] else "Not Found"
 
-    result["date"] = (
-        result["date"][0]
-        if result["date"] else "Not Found"
-    )
+        # CONFIDENCE
+        company_conf = avg_conf(conf_store["company"])
+        date_conf = avg_conf(conf_store["date"])
+        total_conf = avg_conf(conf_store["total"])
 
-    result["total"] = (
-        result["total"][-1]
-        if result["total"] else "Not Found"
-    )
+        overall = (company_conf + date_conf + total_conf) / 3
 
-    return result
+        return {
+            "company": company,
+            "date": date,
+            "total": total,
+            "confidence": round(overall, 3)
+        }
+
+    except Exception as e:
+        return {"error": str(e)}
 
 # =====================================================
-# SEND EMAIL
+# DECISION ENGINE
 # =====================================================
-def send_claim_email(to_email, extracted):
+def decision_layer(conf):
 
-    if not EMAIL_USER or not EMAIL_PASS:
-        return "Email secrets not configured."
+    if conf >= 0.80:
+        return "AUTO_SEND"
 
-    subject = "Insurance Claim Request"
+    elif conf >= 0.60:
+        return "REVIEW"
 
-    body = f"""
-Dear Claims Team,
+    else:
+        return "REJECT"
 
-I would like to request reimbursement for an eligible expense.
+# =====================================================
+# EMAIL SEND
+# =====================================================
+def send_claim_email(to_email, extracted):
 
-Provider Name: {extracted['company']}
-Bill Date: {extracted['date']}
-Claim Amount: ₹{extracted['total']}
+    if not RESEND_API_KEY:
+        return "❌ Missing RESEND_API_KEY secret"
 
-Please process the claim.
+    subject = "Insurance Claim Request"
 
-Regards
-Customer
-"""
+    html_body = f"""
+    <h2>Insurance Claim Request</h2>
 
-    msg = MIMEMultipart()
-    msg["From"] = EMAIL_USER
-    msg["To"] = to_email
-    msg["Subject"] = subject
+    <p>Dear Claims Team,</p>
 
-    msg.attach(MIMEText(body, "plain"))
+    <p>Please process reimbursement request.</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>
+    """
 
     try:
-        server = smtplib.SMTP("smtp.gmail.com", 587)
-        server.starttls()
-        server.login(EMAIL_USER, EMAIL_PASS)
-        server.sendmail(
-            EMAIL_USER,
-            to_email,
-            msg.as_string()
+        response = requests.post(
+            "https://api.resend.com/emails",
+            headers={
+                "Authorization": f"Bearer {RESEND_API_KEY}",
+                "Content-Type": "application/json"
+            },
+            json={
+                "from": FROM_EMAIL,
+                "to": [to_email],
+                "subject": subject,
+                "html": html_body
+            },
+            timeout=20
         )
-        server.quit()
 
-        return f"✅ Email sent successfully to {to_email}"
+        if response.status_code in [200, 201]:
+            return f"✅ Email sent successfully to {to_email}"
+
+        return f"❌ Email failed: {response.text}"
 
     except Exception as e:
-        return f"❌ Email failed: {str(e)}"
+        return f"❌ Email error: {str(e)}"
 
 # =====================================================
-# MAIN UI FUNCTION
+# MAIN PIPELINE
 # =====================================================
 def process_and_send(image, email_id):
 
@@ -190,7 +265,19 @@ def process_and_send(image, email_id):
     if "error" in extracted:
         return extracted, extracted["error"]
 
-    email_status = send_claim_email(email_id, extracted)
+    conf = extracted["confidence"]
+    decision = decision_layer(conf)
+
+    extracted["decision"] = decision
+
+    if decision == "AUTO_SEND":
+        email_status = send_claim_email(email_id, extracted)
+
+    elif decision == "REVIEW":
+        email_status = f"⚠️ Human review required (confidence={conf})"
+
+    else:
+        email_status = f"❌ Rejected (low confidence={conf})"
 
     return extracted, email_status
 
@@ -199,16 +286,19 @@ def process_and_send(image, email_id):
 # =====================================================
 demo = gr.Interface(
     fn=process_and_send,
+
     inputs=[
         gr.Image(type="pil", label="Upload Receipt"),
-        gr.Textbox(label="Insurance Email ID")
+        gr.Textbox(label="Enter Destination Email")
     ],
+
     outputs=[
-        gr.JSON(label="Extracted Data"),
+        gr.JSON(label="AI Extraction"),
         gr.Textbox(label="Email Status")
     ],
+
     title="📄 AI Insurance Claim Generator",
-    description="Upload receipt → Extract details → Auto send claim email"
+    description="Upload receipt → Better extraction → Confidence check → Auto Email"
 )
 
 demo.launch()