Update app.py

app.py

@@ -5,6 +5,8 @@ import torch
 import re
 import requests
 import os
+import base64
+import tempfile
 
 from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 
@@ -13,21 +15,12 @@ from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 # =====================================================
 RESEND_API_KEY = os.getenv("RESEND_API_KEY")
 
-# 🔥 IMPORTANT:
-# Use YOUR VERIFIED DOMAIN EMAIL
-# Example:
-# claims@send.yudham.com
+# Use your VERIFIED sender email
 FROM_EMAIL = "AI Claims <claims@yudham.com>"
 
 MODEL_NAME = "ngupta2026/sroie-layoutlm"
 
-label2id = {
-    "O": 0,
-    "COMPANY": 1,
-    "DATE": 2,
-    "TOTAL": 3
-}
-
+label2id = {"O": 0, "COMPANY": 1, "DATE": 2, "TOTAL": 3}
 id2label = {v: k for k, v in label2id.items()}
 
 # =====================================================
@@ -42,7 +35,7 @@ model.to(device)
 model.eval()
 
 # =====================================================
-# NORMALIZE BOX
+# NORMALIZE BOXES
 # =====================================================
 def normalize(box, width, height):
     return [
@@ -53,7 +46,7 @@ def normalize(box, width, height):
     ]
 
 # =====================================================
-# CONFIDENCE AVG
+# CONFIDENCE HELPER
 # =====================================================
 def avg_conf(lst):
     if len(lst) == 0:
@@ -151,27 +144,15 @@ def extract_receipt(image):
                 except:
                     pass
 
-        # FINAL CLEAN
-        result["company"] = (
-            " ".join(result["company"])
-            if result["company"] else "Not Found"
-        )
-
-        result["date"] = (
-            result["date"][0]
-            if result["date"] else "Not Found"
-        )
-
-        result["total"] = (
-            result["total"][-1]
-            if result["total"] else "Not Found"
-        )
-
-        company_conf = avg_conf(conf_store["company"])
-        date_conf = avg_conf(conf_store["date"])
-        total_conf = avg_conf(conf_store["total"])
+        result["company"] = " ".join(result["company"]) if result["company"] else "Not Found"
+        result["date"] = result["date"][0] if result["date"] else "Not Found"
+        result["total"] = result["total"][-1] if result["total"] else "Not Found"
 
-        overall = (company_conf + date_conf + total_conf) / 3
+        overall = (
+            avg_conf(conf_store["company"]) +
+            avg_conf(conf_store["date"]) +
+            avg_conf(conf_store["total"])
+        ) / 3
 
         result["confidence"] = round(overall, 3)
 
@@ -181,7 +162,7 @@ def extract_receipt(image):
         return {"error": str(e)}
 
 # =====================================================
-# DECISION ENGINE
+# DECISION LAYER
 # =====================================================
 def decision_layer(conf):
 
@@ -195,12 +176,33 @@ def decision_layer(conf):
         return "REJECT"
 
 # =====================================================
-# EMAIL SEND
+# IMAGE TO PDF BASE64
+# =====================================================
+def create_pdf_base64(image):
+
+    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
+
+        pdf_path = f.name
+
+    image = image.convert("RGB")
+    image.save(pdf_path, "PDF")
+
+    with open(pdf_path, "rb") as file:
+        pdf_bytes = file.read()
+
+    encoded = base64.b64encode(pdf_bytes).decode("utf-8")
+
+    return encoded
+
+# =====================================================
+# SEND EMAIL WITH PDF ATTACHMENT
 # =====================================================
-def send_claim_email(to_email, extracted):
+def send_claim_email(to_email, extracted, image):
 
     if not RESEND_API_KEY:
-        return "❌ Missing RESEND_API_KEY secret"
+        return "❌ Missing RESEND_API_KEY"
+
+    pdf_base64 = create_pdf_base64(image)
 
     subject = "Insurance Claim Request"
 
@@ -209,15 +211,30 @@ def send_claim_email(to_email, extracted):
 
     <p>Dear Claims Team,</p>
 
-    <p>Please process reimbursement request.</p>
+    <p>Please process claim reimbursement request.</p>
 
-    <p><b>Provider Name:</b> {extracted['company']}</p>
+    <p><b>Provider:</b> {extracted['company']}</p>
     <p><b>Bill Date:</b> {extracted['date']}</p>
     <p><b>Claim Amount:</b> ₹{extracted['total']}</p>
 
+    <p>Attached: Receipt PDF</p>
+
     <p>Regards,<br>AI Claims System</p>
     """
 
+    payload = {
+        "from": FROM_EMAIL,
+        "to": [to_email],
+        "subject": subject,
+        "html": html_body,
+        "attachments": [
+            {
+                "filename": "receipt.pdf",
+                "content": pdf_base64
+            }
+        ]
+    }
+
     try:
         response = requests.post(
             "https://api.resend.com/emails",
@@ -225,17 +242,12 @@ def send_claim_email(to_email, extracted):
                 "Authorization": f"Bearer {RESEND_API_KEY}",
                 "Content-Type": "application/json"
             },
-            json={
-                "from": FROM_EMAIL,
-                "to": [to_email],
-                "subject": subject,
-                "html": html_body
-            },
-            timeout=20
+            json=payload,
+            timeout=30
         )
 
         if response.status_code in [200, 201]:
-            return f"✅ Email sent successfully to {to_email}"
+            return f"✅ Email + PDF sent to {to_email}"
 
         else:
             return f"❌ Email failed: {response.text}"
@@ -259,15 +271,15 @@ def process_and_send(image, email_id):
     extracted["decision"] = decision
 
     if decision == "AUTO_SEND":
-        email_status = send_claim_email(email_id, extracted)
+        status = send_claim_email(email_id, extracted, image)
 
     elif decision == "REVIEW":
-        email_status = f"⚠️ Human review required (confidence={conf})"
+        status = f"⚠️ Human review needed (confidence={conf})"
 
     else:
-        email_status = f"❌ Rejected (low confidence={conf})"
+        status = f"❌ Rejected (confidence={conf})"
 
-    return extracted, email_status
+    return extracted, status
 
 # =====================================================
 # UI
@@ -281,12 +293,12 @@ demo = gr.Interface(
     ],
 
     outputs=[
-        gr.JSON(label="AI Extraction"),
+        gr.JSON(label="Extracted Fields"),
         gr.Textbox(label="Email Status")
     ],
 
-    title="📄 AI Insurance Claim Generator",
-    description="Upload receipt → Extract fields → Confidence Check → Auto Email"
+    title="📄 AI Insurance Claim Generator + PDF",
+    description="Upload receipt → Extract data → Auto email with PDF attachment"
 )
 
 demo.launch()