Update app.py

app.py

@@ -1,61 +1,73 @@
-import pytesseract
 from PIL import Image
 import torch
+import re
+import os
+import smtplib
 
-from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
+from email.mime.text import MIMEText
+from email.mime.multipart import MIMEMultipart
 
+from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 
+# =====================================================
 # LABELS
+# =====================================================
+label2id = {"O": 0, "COMPANY": 1, "DATE": 2, "TOTAL": 3}
+id2label = {v: k for k, v in label2id.items()}
 
-label2id = {"O":0, "COMPANY":1, "DATE":2, "TOTAL":3}
-id2label = {v:k for k,v in label2id.items()}
-
-# ✅ USE YOUR MODEL
+# =====================================================
+# LOAD MODEL
+# =====================================================
 MODEL_NAME = "ngupta2026/sroie-layoutlm"
 
-
-
 model = LayoutLMForTokenClassification.from_pretrained(MODEL_NAME)
 tokenizer = LayoutLMTokenizerFast.from_pretrained(MODEL_NAME)
-
-# NORMALIZE
-
-
-
-
-
-
+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
+# =====================================================
 def normalize(box, width, height):
     return [
         int(1000 * box[0] / width),
         int(1000 * box[3] / height),
     ]
 
+# =====================================================
+# EXTRACT DATA
+# =====================================================
+def extract_receipt(image):
 
-# MAIN FUNCTION
-
-def process(image):
-
-
-    data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
+    data = pytesseract.image_to_data(
+        image,
+        output_type=pytesseract.Output.DICT
+    )
 
     words = []
     boxes = []
 
     for i in range(len(data["text"])):
-        if data["text"][i].strip() != "":
 
+        text = data["text"][i].strip()
 
-            x = data["left"][i]
-            y = data["top"][i]
-            w = data["width"][i]
+        if text != "":
             h = data["height"][i]
 
-            words.append(data["text"][i])
-            boxes.append([x, y, x+w, y+h])
-
-
+            words.append(text)
+            boxes.append([x, y, x + w, y + h])
 
+    if len(words) == 0:
+        return {"error": "No text detected"}
 
 
     width, height = image.size
@@ -65,13 +77,10 @@ def process(image):
     encoding = tokenizer(
         words,
         boxes=boxes,
-        padding="max_length",
-        truncation=True,
-        is_split_into_words=True,
-        max_length=128
+        max_length=512
     )
 
-
+    encoding = {k: v.to(device) for k, v in encoding.items()}
 
 
     with torch.no_grad():
@@ -79,50 +88,127 @@ def process(image):
 
     predictions = torch.argmax(outputs.logits, dim=2)[0][:len(words)]
 
-    result = {}
-
-
-
+    result = {
+        "company": [],
+        "date": [],
+        "total": []
+    }
 
     for word, pred in zip(words, predictions):
 
         label = id2label[pred.item()]
 
-
+        # company from model
         if label == "COMPANY":
-            result.setdefault("company", []).append(word)
-        elif label == "DATE":
-            result.setdefault("date", []).append(word)
-        elif label == "TOTAL":
-            result.setdefault("total", []).append(word)
+            result["company"].append(word)
+
+        # date from regex
+        if re.search(r"\d{2}[/-]\d{2}[/-]\d{2,4}", word):
+            result["date"].append(word)
+
+        # 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
+
+    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"
+    )
 
+    return result
 
+# =====================================================
+# SEND EMAIL
+# =====================================================
+def send_claim_email(to_email, extracted):
 
+    if not EMAIL_USER or not EMAIL_PASS:
+        return "Email secrets not configured."
 
+    subject = "Insurance Claim Request"
 
+    body = f"""
+Dear Claims Team,
 
+I would like to request reimbursement for an eligible expense.
 
+Provider Name: {extracted['company']}
+Bill Date: {extracted['date']}
+Claim Amount: ₹{extracted['total']}
 
+Please process the claim.
 
+Regards
+Customer
+"""
 
+    msg = MIMEMultipart()
+    msg["From"] = EMAIL_USER
+    msg["To"] = to_email
+    msg["Subject"] = subject
 
+    msg.attach(MIMEText(body, "plain"))
 
+    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()
+        )
+        server.quit()
 
+        return f"✅ Email sent successfully to {to_email}"
 
+    except Exception as e:
+        return f"❌ Email failed: {str(e)}"
 
+# =====================================================
+# MAIN UI FUNCTION
+# =====================================================
+def process_and_send(image, email_id):
 
-    return result
+    extracted = extract_receipt(image)
 
+    if "error" in extracted:
+        return extracted, extracted["error"]
 
-# UI
+    email_status = send_claim_email(email_id, extracted)
 
-demo = gr.Interface(
-    fn=process,
-    inputs=gr.Image(type="pil"),
-    outputs="json",
-    title="Document AI Extractor"
+    return extracted, email_status
 
+# =====================================================
+# UI
+# =====================================================
+demo = gr.Interface(
+    fn=process_and_send,
+    inputs=[
+        gr.Image(type="pil", label="Upload Receipt"),
+        gr.Textbox(label="Insurance Email ID")
+    ],
+    outputs=[
+        gr.JSON(label="Extracted Data"),
+        gr.Textbox(label="Email Status")
+    ],
+    title="📄 AI Insurance Claim Generator",
+    description="Upload receipt → Extract details → Auto send claim email"
 )
 
 demo.launch()