Update app.py

app.py

@@ -1,3 +1,10 @@
+# =====================================================
+# AI INSURANCE CLAIM GENERATOR
+# FINAL VERSION
+# Working Extraction Logic + Professional PDF + Email
+# Hugging Face Space Ready
+# =====================================================
+
 import gradio as gr
 import pytesseract
 from PIL import Image
@@ -10,7 +17,7 @@ import base64
 
 from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 
-# PDF
+# PDF LIBRARY
 from reportlab.lib.pagesizes import A4
 from reportlab.pdfgen import canvas
 
@@ -19,9 +26,7 @@ from reportlab.pdfgen import canvas
 # =====================================================
 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 +50,7 @@ model.to(device)
 model.eval()
 
 # =====================================================
-# NORMALIZE BOXES
+# HELPERS
 # =====================================================
 def normalize(box, width, height):
     return [
@@ -55,21 +60,76 @@ def normalize(box, width, height):
         int(1000 * box[3] / height),
     ]
 
+def avg(lst):
+    return sum(lst) / len(lst) if len(lst) > 0 else 0
+
 # =====================================================
-# HELPERS
+# CLEAN COMPANY
 # =====================================================
-def clean_text(txt):
-    return txt.strip().replace("\n", " ")
+def clean_company(txt):
 
-def avg(lst):
-    if len(lst) == 0:
-        return 0
-    return sum(lst) / len(lst)
+    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"
+
+    return txt.upper()
+
+# =====================================================
+# DATE EXTRACTION
+# =====================================================
+def extract_date(words):
+
+    for w in words:
+        if re.fullmatch(r"\d{1,2}[/-]\d{1,2}[/-]\d{2,4}", w):
+            return w
+
+    return "Not Found"
 
 # =====================================================
-# PDF CREATION
+# TOTAL EXTRACTION
 # =====================================================
-def create_pdf(extracted):
+def clean_amount_token(txt):
+
+    txt = txt.upper()
+    txt = txt.replace("RM", "")
+    txt = txt.replace("MYR", "")
+    txt = txt.replace("RS", "")
+    txt = txt.replace("₹", "")
+    txt = txt.replace(",", "")
+    txt = txt.strip()
+
+    return txt
+
+def extract_total(words):
+
+    vals = []
+
+    for w in words:
+
+        x = clean_amount_token(w)
+
+        if re.fullmatch(r"\d+\.\d{2}", x):
+            try:
+                v = float(x)
+
+                if 0.5 <= v <= 100000:
+                    vals.append(v)
+
+            except:
+                pass
+
+    if vals:
+        return f"{max(vals):.2f}"
+
+    return "Not Found"
+
+# =====================================================
+# PROFESSIONAL PDF GENERATOR
+# =====================================================
+def create_pdf_base64(extracted):
 
     buffer = io.BytesIO()
 
@@ -79,7 +139,7 @@ def create_pdf(extracted):
     y = height - 60
 
     p.setFont("Helvetica-Bold", 18)
-    p.drawString(50, y, "Insurance Claim Report")
+    p.drawString(50, y, "AI Insurance Claim Report")
 
     y -= 40
     p.setFont("Helvetica", 12)
@@ -87,17 +147,17 @@ def create_pdf(extracted):
     rows = [
         f"Provider Name : {extracted['company']}",
         f"Bill Date     : {extracted['date']}",
-        f"Claim Amount  : Rs {extracted['total']}",
+        f"Claim Amount  : {extracted['total']}",
         f"Confidence    : {extracted['confidence']}",
-        f"Decision      : {extracted['decision']}",
+        f"Decision      : {extracted['decision']}"
     ]
 
     for row in rows:
         p.drawString(50, y, row)
-        y -= 25
+        y -= 28
 
     y -= 20
-    p.drawString(50, y, "Generated by AI Insurance Claim System")
+    p.drawString(50, y, "Generated by AI Insurance Claim Generator")
 
     p.showPage()
     p.save()
@@ -105,15 +165,16 @@ def create_pdf(extracted):
     pdf_bytes = buffer.getvalue()
     buffer.close()
 
-    return base64.b64encode(pdf_bytes).decode("utf-8")
+    return base64.b64encode(pdf_bytes).decode()
 
 # =====================================================
-# EXTRACTION ENGINE
+# MAIN EXTRACTION
 # =====================================================
 def extract_receipt(image):
 
     try:
         image = image.convert("RGB")
+        image.thumbnail((1500, 1500))
 
         data = pytesseract.image_to_data(
             image,
@@ -122,203 +183,148 @@ 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 found"}
+            return {"error": "No text detected"}
 
         width, height = image.size
-        boxes_norm = [normalize(b, width, height) for b in boxes]
+        boxes = [normalize(b, width, height) for b in boxes]
 
-        # =========================
-        # TOKENIZER
-        # =========================
         encoding = tokenizer(
             words,
-            boxes=boxes_norm,
+            boxes=boxes,
             return_tensors="pt",
             truncation=True,
             padding="max_length",
-            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
-        # =========================
         with torch.no_grad():
             outputs = model(**encoding)
 
         probs = torch.softmax(outputs.logits, dim=2)
+
         preds = torch.argmax(probs, dim=2)[0][:len(words)]
-        pred_conf = torch.max(probs, dim=2)[0][0][:len(words)]
+        confs = torch.max(probs, dim=2)[0][0][:len(words)]
 
-        company_words = []
+        company_tokens = []
         company_scores = []
 
-        # =========================
-        # COMPANY FROM MODEL
-        # =========================
-        for word, pred, c in zip(words, preds, pred_conf):
+        for word, pred, conf in zip(words, preds, confs):
 
             label = id2label[pred.item()]
 
             if label == "COMPANY":
-                company_words.append(word)
-                company_scores.append(c.item())
-
-        company = " ".join(company_words).strip()
-
-        if company == "":
-            company = words[0]
-
-        # =========================
-        # DATE BY REGEX
-        # =========================
-        date = "Not Found"
+                company_tokens.append(word)
+                company_scores.append(conf.item())
 
-        for w in words:
-            if re.search(r"\d{2}[/-]\d{2}[/-]\d{2,4}", w):
-                date = w
-                break
-
-        # =========================
-        # 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 company_tokens:
+            company = " ".join(company_tokens[:8])
+        else:
+            company = " ".join(words[:5])
 
-                    if 1 <= val <= 100000:
-                        amount_candidates.append(val)
-                except:
-                    pass
+        company = clean_company(company)
 
-        total = "Not Found"
+        date = extract_date(words)
+        total = extract_total(words)
 
-        if len(amount_candidates) > 0:
-            total = f"{max(amount_candidates):.2f}"
+        score = avg(company_scores)
 
-        # =========================
-        # CONFIDENCE
-        # =========================
-        company_conf = avg(company_scores)
+        if date != "Not Found":
+            score += 0.12
 
         if total != "Not Found":
-            total_conf = 0.90
-        else:
-            total_conf = 0.20
+            score += 0.18
 
-        if date != "Not Found":
-            date_conf = 0.90
-        else:
-            date_conf = 0.20
+        score = min(score, 0.99)
 
-        overall = round((company_conf + total_conf + date_conf) / 3, 3)
-
-        result = {
+        return {
             "company": company,
             "date": date,
             "total": total,
-            "confidence": overall
+            "confidence": round(score, 3)
         }
 
-        return result
-
     except Exception as e:
         return {"error": str(e)}
 
 # =====================================================
-# DECISION ENGINE
+# DECISION
 # =====================================================
 def decision_layer(conf):
 
     if conf >= 0.80:
         return "AUTO_SEND"
-
     elif conf >= 0.60:
         return "REVIEW"
-
-    return "REJECT"
+    else:
+        return "REJECT"
 
 # =====================================================
-# EMAIL SEND WITH PDF
+# EMAIL SEND
 # =====================================================
 def send_claim_email(to_email, extracted):
 
     if not RESEND_API_KEY:
         return "❌ Missing RESEND_API_KEY"
 
-    pdf_base64 = create_pdf(extracted)
+    pdf_b64 = create_pdf_base64(extracted)
 
-    html_body = f"""
-    <h2>Insurance Claim Request</h2>
-
-    <p>Dear Claims Team,</p>
+    subject = "Insurance Claim Request"
 
-    <p>Please process reimbursement request.</p>
+    html = f"""
+    <h2>Insurance Claim Request</h2>
 
-    <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><b>Provider:</b> {extracted['company']}</p>
+    <p><b>Date:</b> {extracted['date']}</p>
+    <p><b>Amount:</b> ₹{extracted['total']}</p>
 
-    <p>Regards,<br>AI Claims System</p>
+    <p>Attached: AI Claim Report PDF</p>
     """
 
-    payload = {
-        "from": FROM_EMAIL,
-        "to": [to_email],
-        "subject": "Insurance Claim Request",
-        "html": html_body,
-        "attachments": [
-            {
-                "filename": "Claim_Report.pdf",
-                "content": pdf_base64
-            }
-        ]
-    }
-
     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,
+            json={
+                "from": FROM_EMAIL,
+                "to": [to_email],
+                "subject": subject,
+                "html": html,
+                "attachments": [
+                    {
+                        "filename": "claim_report.pdf",
+                        "content": pdf_b64
+                    }
+                ]
+            },
             timeout=20
         )
 
-        if response.status_code in [200, 201]:
-            return f"✅ Email sent successfully to {to_email}"
+        if r.status_code in [200, 201]:
+            return f"✅ Email + PDF sent to {to_email}"
 
-        return f"❌ Email failed: {response.text}"
+        return f"❌ Email failed: {r.text}"
 
     except Exception as e:
         return f"❌ Email error: {str(e)}"
@@ -333,19 +339,21 @@ def process_and_send(image, email_id):
     if "error" in extracted:
         return extracted, extracted["error"]
 
-    decision = decision_layer(extracted["confidence"])
+    conf = extracted["confidence"]
+    decision = decision_layer(conf)
+
     extracted["decision"] = decision
 
     if decision == "AUTO_SEND":
-        email_status = send_claim_email(email_id, extracted)
+        status = send_claim_email(email_id, extracted)
 
     elif decision == "REVIEW":
-        email_status = f"⚠️ Needs manual review (confidence={extracted['confidence']})"
+        status = f"⚠️ Human review required ({conf})"
 
     else:
-        email_status = f"❌ Rejected (confidence={extracted['confidence']})"
+        status = f"❌ Rejected ({conf})"
 
-    return extracted, email_status
+    return extracted, status
 
 # =====================================================
 # UI
@@ -355,16 +363,16 @@ demo = gr.Interface(
 
     inputs=[
         gr.Image(type="pil", label="Upload Receipt"),
-        gr.Textbox(label="Enter Destination Email")
+        gr.Textbox(label="Destination Email")
     ],
 
     outputs=[
-        gr.JSON(label="AI Extraction Result"),
+        gr.JSON(label="AI Extraction"),
         gr.Textbox(label="Email Status")
     ],
 
     title="📄 AI Insurance Claim Generator",
-    description="Upload receipt → Extract fields → Confidence Check → Auto Email + PDF"
+    description="Upload receipt → Extract fields → Generate PDF → Auto Email"
 )
 
 demo.launch()