Update app.py

app.py

@@ -1,9 +1,3 @@
-# =====================================================
-# AI INSURANCE CLAIM GENERATOR (FINAL + PDF VERSION)
-# Accurate Extraction + PDF + Email
-# Hugging Face Space Ready
-# =====================================================
-
 import gradio as gr
 import pytesseract
 from PIL import Image
@@ -16,12 +10,18 @@ import base64
 
 from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 
+# PDF
+from reportlab.lib.pagesizes import A4
+from reportlab.pdfgen import canvas
+
 # =====================================================
 # CONFIG
 # =====================================================
 RESEND_API_KEY = os.getenv("RESEND_API_KEY")
 
+# Use your verified sender email/domain
 FROM_EMAIL = "AI Claims <claims@yudham.com>"
+
 MODEL_NAME = "ngupta2026/sroie-layoutlm"
 
 label2id = {
@@ -45,7 +45,7 @@ model.to(device)
 model.eval()
 
 # =====================================================
-# HELPERS
+# NORMALIZE BOXES
 # =====================================================
 def normalize(box, width, height):
     return [
@@ -55,78 +55,65 @@ def normalize(box, width, height):
         int(1000 * box[3] / height),
     ]
 
-def avg(lst):
-    return sum(lst) / len(lst) if len(lst) > 0 else 0
-
 # =====================================================
-# COMPANY
+# HELPERS
 # =====================================================
-def clean_company(txt):
-
-    txt = txt.strip()
-    txt = re.sub(r"[^A-Za-z0-9&().,\- /]", "", txt)
-    txt = re.sub(r"\s+", " ", txt).strip()
-
-    if len(txt) < 2:
-        return "Not Found"
+def clean_text(txt):
+    return txt.strip().replace("\n", " ")
 
-    return txt.upper()
+def avg(lst):
+    if len(lst) == 0:
+        return 0
+    return sum(lst) / len(lst)
 
 # =====================================================
-# DATE
+# PDF CREATION
 # =====================================================
-def extract_date(words):
-
-    for w in words:
-        if re.fullmatch(r"\d{1,2}[/-]\d{1,2}[/-]\d{2,4}", w):
-            return w
+def create_pdf(extracted):
 
-    return "Not Found"
+    buffer = io.BytesIO()
 
-# =====================================================
-# TOTAL
-# =====================================================
-def clean_amount_token(txt):
+    p = canvas.Canvas(buffer, pagesize=A4)
+    width, height = A4
 
-    txt = txt.upper()
-    txt = txt.replace("RM", "")
-    txt = txt.replace("MYR", "")
-    txt = txt.replace("RS", "")
-    txt = txt.replace("₹", "")
-    txt = txt.replace(",", "")
-    txt = txt.strip()
+    y = height - 60
 
-    return txt
+    p.setFont("Helvetica-Bold", 18)
+    p.drawString(50, y, "Insurance Claim Report")
 
-def extract_total(words):
+    y -= 40
+    p.setFont("Helvetica", 12)
 
-    vals = []
+    rows = [
+        f"Provider Name : {extracted['company']}",
+        f"Bill Date     : {extracted['date']}",
+        f"Claim Amount  : Rs {extracted['total']}",
+        f"Confidence    : {extracted['confidence']}",
+        f"Decision      : {extracted['decision']}",
+    ]
 
-    for w in words:
+    for row in rows:
+        p.drawString(50, y, row)
+        y -= 25
 
-        x = clean_amount_token(w)
+    y -= 20
+    p.drawString(50, y, "Generated by AI Insurance Claim System")
 
-        if re.fullmatch(r"\d+\.\d{2}", x):
-            try:
-                v = float(x)
-                if 0.5 <= v <= 100000:
-                    vals.append(v)
-            except:
-                pass
+    p.showPage()
+    p.save()
 
-    if vals:
-        return f"{max(vals):.2f}"
+    pdf_bytes = buffer.getvalue()
+    buffer.close()
 
-    return "Not Found"
+    return base64.b64encode(pdf_bytes).decode("utf-8")
 
 # =====================================================
-# EXTRACTION
+# EXTRACTION ENGINE
 # =====================================================
 def extract_receipt(image):
 
     try:
         image = image.convert("RGB")
-        image.thumbnail((1500, 1500))
 
         data = pytesseract.image_to_data(
             image,
@@ -135,200 +122,209 @@ def extract_receipt(image):
 
         words = []
         boxes = []
+        confs = []
+
+        n = len(data["text"])
 
-        for i in range(len(data["text"])):
+        for i in range(n):
 
-            txt = data["text"][i].strip()
+            txt = clean_text(data["text"][i])
 
-            if txt != "" and len(txt) > 1:
+            if txt == "":
+                continue
 
-                x = data["left"][i]
-                y = data["top"][i]
-                w = data["width"][i]
-                h = data["height"][i]
+            x = data["left"][i]
+            y = data["top"][i]
+            w = data["width"][i]
+            h = data["height"][i]
 
-                words.append(txt)
-                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_norm = [normalize(b, width, height) for b in boxes]
 
+        # =========================
+        # TOKENIZER
+        # =========================
         encoding = tokenizer(
             words,
-            boxes=boxes,
+            boxes=boxes_norm,
             return_tensors="pt",
             truncation=True,
             padding="max_length",
-            max_length=512,
-            is_split_into_words=True
+            is_split_into_words=True,
+            max_length=256
         )
 
         encoding = {k: v.to(device) for k, v in encoding.items()}
 
+        # =========================
+        # MODEL
+        # =========================
         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)]
+        pred_conf = torch.max(probs, dim=2)[0][0][:len(words)]
 
-        company_tokens = []
+        company_words = []
         company_scores = []
 
-        for word, pred, conf in zip(words, preds, confs):
+        # =========================
+        # COMPANY FROM MODEL
+        # =========================
+        for word, pred, c in zip(words, preds, pred_conf):
 
             label = id2label[pred.item()]
 
             if label == "COMPANY":
-                company_tokens.append(word)
-                company_scores.append(conf.item())
+                company_words.append(word)
+                company_scores.append(c.item())
 
-        if company_tokens:
-            company = " ".join(company_tokens[:8])
-        else:
-            company = " ".join(words[:5])
+        company = " ".join(company_words).strip()
 
-        company = clean_company(company)
+        if company == "":
+            company = words[0]
 
-        date = extract_date(words)
-        total = extract_total(words)
+        # =========================
+        # DATE BY REGEX
+        # =========================
+        date = "Not Found"
 
-        score = avg(company_scores)
+        for w in words:
+            if re.search(r"\d{2}[/-]\d{2}[/-]\d{2,4}", w):
+                date = w
+                break
 
-        if date != "Not Found":
-            score += 0.12
+        # =========================
+        # TOTAL SMART LOGIC
+        # =========================
+        amount_candidates = []
+
+        for w in words:
+
+            t = w.replace(",", "").replace("RM", "").replace("Rs", "").replace("₹", "")
+
+            if re.fullmatch(r"\d+(\.\d{2})?", t):
+                try:
+                    val = float(t)
+
+                    if 1 <= val <= 100000:
+                        amount_candidates.append(val)
+                except:
+                    pass
+
+        total = "Not Found"
+
+        if len(amount_candidates) > 0:
+            total = f"{max(amount_candidates):.2f}"
+
+        # =========================
+        # CONFIDENCE
+        # =========================
+        company_conf = avg(company_scores)
 
         if total != "Not Found":
-            score += 0.18
+            total_conf = 0.90
+        else:
+            total_conf = 0.20
 
-        score = min(score, 0.99)
+        if date != "Not Found":
+            date_conf = 0.90
+        else:
+            date_conf = 0.20
+
+        overall = round((company_conf + total_conf + date_conf) / 3, 3)
 
-        return {
+        result = {
             "company": company,
             "date": date,
             "total": total,
-            "confidence": round(score, 3)
+            "confidence": overall
         }
 
+        return result
+
     except Exception as e:
         return {"error": str(e)}
 
 # =====================================================
-# DECISION
+# DECISION ENGINE
 # =====================================================
 def decision_layer(conf):
 
     if conf >= 0.80:
         return "AUTO_SEND"
+
     elif conf >= 0.60:
         return "REVIEW"
-    else:
-        return "REJECT"
 
-# =====================================================
-# SIMPLE PDF GENERATOR (NO EXTRA LIBRARY NEEDED)
-# =====================================================
-def create_pdf_bytes(extracted):
-
-    text = f"""
-AI INSURANCE CLAIM REPORT
-
-Provider Name : {extracted['company']}
-Bill Date     : {extracted['date']}
-Claim Amount  : {extracted['total']}
-Confidence    : {extracted['confidence']}
-Decision      : {extracted['decision']}
-
-Generated by AI Claims System
-"""
-
-    # Minimal PDF binary
-    pdf = f"""%PDF-1.4
-1 0 obj<<>>endobj
-2 0 obj<< /Length {len(text)+80} >>stream
-BT
-/F1 12 Tf
-50 750 Td
-({text.replace(chr(10), ') Tj T* (')}) Tj
-ET
-endstream
-endobj
-3 0 obj<< /Type /Page /Parent 4 0 R /Contents 2 0 R >>endobj
-4 0 obj<< /Type /Pages /Kids [3 0 R] /Count 1 >>endobj
-5 0 obj<< /Type /Catalog /Pages 4 0 R >>endobj
-xref
-0 6
-0000000000 65535 f
-0000000010 00000 n
-0000000030 00000 n
-0000000000 00000 n
-0000000000 00000 n
-0000000000 00000 n
-trailer<< /Root 5 0 R /Size 6 >>
-startxref
-0
-%%EOF
-"""
-    return pdf.encode("latin-1", errors="ignore")
+    return "REJECT"
 
 # =====================================================
-# EMAIL WITH PDF ATTACHMENT
+# EMAIL SEND WITH PDF
 # =====================================================
 def send_claim_email(to_email, extracted):
 
     if not RESEND_API_KEY:
         return "❌ Missing RESEND_API_KEY"
 
-    pdf_bytes = create_pdf_bytes(extracted)
-    pdf_b64 = base64.b64encode(pdf_bytes).decode()
-
-    subject = "Insurance Claim Request"
+    pdf_base64 = create_pdf(extracted)
 
-    html = f"""
+    html_body = f"""
     <h2>Insurance Claim Request</h2>
 
-    <p><b>Provider:</b> {extracted['company']}</p>
-    <p><b>Date:</b> {extracted['date']}</p>
-    <p><b>Amount:</b> ₹{extracted['total']}</p>
+    <p>Dear Claims Team,</p>
 
-    <p>Attached: AI Claim Report PDF</p>
+    <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> Rs {extracted['total']}</p>
+    <p><b>Confidence:</b> {extracted['confidence']}</p>
+
+    <p>Regards,<br>AI Claims System</p>
     """
 
+    payload = {
+        "from": FROM_EMAIL,
+        "to": [to_email],
+        "subject": "Insurance Claim Request",
+        "html": html_body,
+        "attachments": [
+            {
+                "filename": "Claim_Report.pdf",
+                "content": pdf_base64
+            }
+        ]
+    }
+
     try:
-        r = requests.post(
+        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,
-                "attachments": [
-                    {
-                        "filename": "claim_report.pdf",
-                        "content": pdf_b64
-                    }
-                ]
-            },
+            json=payload,
             timeout=20
         )
 
-        if r.status_code in [200, 201]:
-            return f"✅ Email + PDF sent to {to_email}"
+        if response.status_code in [200, 201]:
+            return f"✅ Email sent successfully to {to_email}"
 
-        return f"❌ Email failed: {r.text}"
+        return f"❌ Email failed: {response.text}"
 
     except Exception as e:
         return f"❌ Email error: {str(e)}"
 
 # =====================================================
-# MAIN
+# MAIN PIPELINE
 # =====================================================
 def process_and_send(image, email_id):
 
@@ -337,37 +333,38 @@ def process_and_send(image, email_id):
     if "error" in extracted:
         return extracted, extracted["error"]
 
-    conf = extracted["confidence"]
-    decision = decision_layer(conf)
-
+    decision = decision_layer(extracted["confidence"])
     extracted["decision"] = decision
 
     if decision == "AUTO_SEND":
-        status = send_claim_email(email_id, extracted)
+        email_status = send_claim_email(email_id, extracted)
 
     elif decision == "REVIEW":
-        status = f"⚠️ Human review required ({conf})"
+        email_status = f"⚠️ Needs manual review (confidence={extracted['confidence']})"
 
     else:
-        status = f"❌ Rejected ({conf})"
+        email_status = f"❌ Rejected (confidence={extracted['confidence']})"
 
-    return extracted, status
+    return extracted, email_status
 
 # =====================================================
 # UI
 # =====================================================
 demo = gr.Interface(
     fn=process_and_send,
+
     inputs=[
         gr.Image(type="pil", label="Upload Receipt"),
-        gr.Textbox(label="Destination Email")
+        gr.Textbox(label="Enter Destination Email")
     ],
+
     outputs=[
-        gr.JSON(label="AI Extraction"),
+        gr.JSON(label="AI Extraction Result"),
         gr.Textbox(label="Email Status")
     ],
+
     title="📄 AI Insurance Claim Generator",
-    description="Upload receipt → Extract fields → Generate PDF → Auto Email"
+    description="Upload receipt → Extract fields → Confidence Check → Auto Email + PDF"
 )
 
 demo.launch()