Update app.py

app.py

@@ -5,8 +5,11 @@ import torch
 import re
 import requests
 import os
+import io
 import base64
-import tempfile
+
+from reportlab.pdfgen import canvas
+from reportlab.lib.pagesizes import A4
 
 from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 
@@ -15,12 +18,18 @@ from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 # =====================================================
 RESEND_API_KEY = os.getenv("RESEND_API_KEY")
 
-# Use your VERIFIED sender email
-FROM_EMAIL = "AI Claims <claims@yudham.com>"
+# VERIFIED DOMAIN EMAIL (CHANGE THIS TO YOUR VERIFIED DOMAIN)
+FROM_EMAIL = "claims@send.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()}
 
 # =====================================================
@@ -46,21 +55,12 @@ def normalize(box, width, height):
     ]
 
 # =====================================================
-# CONFIDENCE HELPER
-# =====================================================
-def avg_conf(lst):
-    if len(lst) == 0:
-        return 0
-    return sum(lst) / len(lst)
-
-# =====================================================
-# OCR + EXTRACTION
+# OCR + MODEL EXTRACTION
 # =====================================================
 def extract_receipt(image):
 
     try:
         image = image.convert("RGB")
-        image.thumbnail((1200, 1200))
 
         data = pytesseract.image_to_data(
             image,
@@ -72,212 +72,216 @@ def extract_receipt(image):
 
         for i in range(len(data["text"])):
 
-            text = data["text"][i].strip()
-
-            if text != "" and len(text) > 1:
+            txt = data["text"][i].strip()
 
+            if txt != "":
                 x = data["left"][i]
                 y = data["top"][i]
                 w = data["width"][i]
                 h = data["height"][i]
 
-                words.append(text)
-                boxes.append([x, y, x + w, y + h])
+                words.append(txt)
+                boxes.append([x, y, x+w, y+h])
 
         if len(words) == 0:
-            return {"error": "No text detected"}
+            return {"error": "No text found"}
 
         width, height = image.size
-        boxes = [normalize(b, width, height) for b in boxes]
+        boxes = [normalize(box, width, height) for box in boxes]
 
+        # =================================================
+        # TOKENIZER
+        # =================================================
         encoding = tokenizer(
             words,
             boxes=boxes,
             return_tensors="pt",
             padding="max_length",
             truncation=True,
-            is_split_into_words=True,
-            max_length=256
+            max_length=512,
+            is_split_into_words=True
         )
 
         encoding = {k: v.to(device) for k, v in encoding.items()}
 
+        # =================================================
+        # MODEL PREDICTION
+        # =================================================
         with torch.no_grad():
             outputs = model(**encoding)
 
         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)]
-
-        result = {
-            "company": [],
-            "date": [],
-            "total": []
-        }
 
-        conf_store = {
-            "company": [],
-            "date": [],
-            "total": []
-        }
+        # =================================================
+        # EXTRACTION STORE
+        # =================================================
+        company_tokens = []
+        totals = []
+        dates = []
 
-        for word, pred, conf in zip(words, preds, confs):
+        for word, pred in zip(words, preds):
 
             label = id2label[pred.item()]
-            c = conf.item()
 
+            # COMPANY
             if label == "COMPANY":
-                result["company"].append(word)
-                conf_store["company"].append(c)
+                company_tokens.append(word)
 
-            if re.search(r"\d{2}[/-]\d{2}[/-]\d{2,4}", word):
-                result["date"].append(word)
-                conf_store["date"].append(c)
+            # DATE via regex
+            if re.search(r"\d{1,2}[/-]\d{1,2}[/-]\d{2,4}", word):
+                dates.append(word)
 
-            if re.search(r"\d+(\.\d{2})?", word):
+            # MONEY
+            if re.search(r"^\d+[.,]?\d*$", word):
                 try:
-                    value = float(word.replace(",", ""))
-                    if value > 50:
-                        result["total"].append(word)
-                        conf_store["total"].append(c)
+                    val = float(word.replace(",", ""))
+                    if val > 20:
+                        totals.append(val)
                 except:
                     pass
 
-        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 = (
-            avg_conf(conf_store["company"]) +
-            avg_conf(conf_store["date"]) +
-            avg_conf(conf_store["total"])
-        ) / 3
-
-        result["confidence"] = round(overall, 3)
-
-        return result
+        # =================================================
+        # FINAL CLEANUP
+        # =================================================
+        company = " ".join(company_tokens[:6]).strip()
+        if company == "":
+            company = "Not Found"
+
+        date = dates[0] if len(dates) > 0 else "Not Found"
+
+        total = str(max(totals)) if len(totals) > 0 else "Not Found"
+
+        # =================================================
+        # ADDRESS HEURISTIC
+        # =================================================
+        address_lines = []
+
+        for w in words:
+            if (
+                w not in company_tokens
+                and w not in dates
+                and not re.search(r"^\d+[.,]?\d*$", w)
+            ):
+                if len(w) > 2:
+                    address_lines.append(w)
+
+        address = " ".join(address_lines[:10]).strip()
+
+        if address == "":
+            address = "Not Found"
+
+        return {
+            "company": company,
+            "date": date,
+            "total": total,
+            "address": address
+        }
 
     except Exception as e:
         return {"error": str(e)}
 
 # =====================================================
-# DECISION LAYER
+# PDF GENERATOR
 # =====================================================
-def decision_layer(conf):
+def create_pdf(extracted):
 
-    if conf >= 0.80:
-        return "AUTO_SEND"
+    buffer = io.BytesIO()
 
-    elif conf >= 0.60:
-        return "REVIEW"
+    c = canvas.Canvas(buffer, pagesize=A4)
+    width, height = A4
 
-    else:
-        return "REJECT"
+    y = height - 60
 
-# =====================================================
-# IMAGE TO PDF BASE64
-# =====================================================
-def create_pdf_base64(image):
+    c.setFont("Helvetica-Bold", 18)
+    c.drawString(50, y, "Insurance Claim Summary")
 
-    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
+    y -= 40
+    c.setFont("Helvetica", 12)
 
-        pdf_path = f.name
+    lines = [
+        f"Provider Name : {extracted['company']}",
+        f"Bill Date     : {extracted['date']}",
+        f"Claim Amount  : ₹{extracted['total']}",
+        f"Address       : {extracted['address']}",
+    ]
 
-    image = image.convert("RGB")
-    image.save(pdf_path, "PDF")
+    for line in lines:
+        c.drawString(50, y, line)
+        y -= 30
 
-    with open(pdf_path, "rb") as file:
-        pdf_bytes = file.read()
+    c.save()
 
-    encoded = base64.b64encode(pdf_bytes).decode("utf-8")
+    pdf_bytes = buffer.getvalue()
+    buffer.close()
 
-    return encoded
+    return pdf_bytes
 
 # =====================================================
-# SEND EMAIL WITH PDF ATTACHMENT
+# EMAIL SEND VIA RESEND
 # =====================================================
-def send_claim_email(to_email, extracted, image):
+def send_email(to_email, extracted):
 
     if not RESEND_API_KEY:
-        return "❌ Missing RESEND_API_KEY"
-
-    pdf_base64 = create_pdf_base64(image)
+        return "❌ RESEND_API_KEY missing"
 
-    subject = "Insurance Claim Request"
+    pdf_data = create_pdf(extracted)
+    pdf_b64 = base64.b64encode(pdf_data).decode()
 
-    html_body = f"""
+    html = f"""
     <h2>Insurance Claim Request</h2>
-
-    <p>Dear Claims Team,</p>
-
-    <p>Please process claim reimbursement request.</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>
+    <p><b>Date:</b> {extracted['date']}</p>
+    <p><b>Amount:</b> ₹{extracted['total']}</p>
+    <p><b>Address:</b> {extracted['address']}</p>
+    <p>Please find attached PDF summary.</p>
     """
 
     payload = {
         "from": FROM_EMAIL,
         "to": [to_email],
-        "subject": subject,
-        "html": html_body,
+        "subject": "Insurance Claim Request",
+        "html": html,
         "attachments": [
             {
-                "filename": "receipt.pdf",
-                "content": pdf_base64
+                "filename": "claim_summary.pdf",
+                "content": pdf_b64
             }
         ]
     }
 
+    headers = {
+        "Authorization": f"Bearer {RESEND_API_KEY}",
+        "Content-Type": "application/json"
+    }
+
     try:
-        response = requests.post(
+        r = requests.post(
             "https://api.resend.com/emails",
-            headers={
-                "Authorization": f"Bearer {RESEND_API_KEY}",
-                "Content-Type": "application/json"
-            },
             json=payload,
-            timeout=30
+            headers=headers,
+            timeout=20
         )
 
-        if response.status_code in [200, 201]:
-            return f"✅ Email + PDF sent to {to_email}"
+        if r.status_code in [200, 201]:
+            return f"✅ Email sent to {to_email}"
 
-        else:
-            return f"❌ Email failed: {response.text}"
+        return f"❌ Email failed: {r.text}"
 
     except Exception as e:
         return f"❌ Email error: {str(e)}"
 
 # =====================================================
-# MAIN PIPELINE
+# MAIN FUNCTION
 # =====================================================
-def process_and_send(image, email_id):
+def process(image, email):
 
     extracted = extract_receipt(image)
 
     if "error" in extracted:
         return extracted, extracted["error"]
 
-    conf = extracted["confidence"]
-    decision = decision_layer(conf)
-
-    extracted["decision"] = decision
-
-    if decision == "AUTO_SEND":
-        status = send_claim_email(email_id, extracted, image)
-
-    elif decision == "REVIEW":
-        status = f"⚠️ Human review needed (confidence={conf})"
-
-    else:
-        status = f"❌ Rejected (confidence={conf})"
+    status = send_email(email, extracted)
 
     return extracted, status
 
@@ -285,20 +289,17 @@ def process_and_send(image, email_id):
 # UI
 # =====================================================
 demo = gr.Interface(
-    fn=process_and_send,
-
+    fn=process,
     inputs=[
         gr.Image(type="pil", label="Upload Receipt"),
-        gr.Textbox(label="Enter Destination Email")
+        gr.Textbox(label="Enter Email ID")
     ],
-
     outputs=[
-        gr.JSON(label="Extracted Fields"),
+        gr.JSON(label="Extracted Output"),
         gr.Textbox(label="Email Status")
     ],
-
-    title="📄 AI Insurance Claim Generator + PDF",
-    description="Upload receipt → Extract data → Auto email with PDF attachment"
+    title="📄 AI Insurance Claim Generator",
+    description="Upload receipt → Extract details → Generate PDF → Send Email"
 )
 
 demo.launch()