Update src/streamlit_app.py

src/streamlit_app.py

@@ -236,25 +236,44 @@ def run_inference_on_image(image: Image.Image, processor, model, device, decoder
 # ---------------------------
 def map_prediction_to_ui(pred):
     import json, re
-    from datetime import datetime
+    from collections import defaultdict
 
+    # --- parse raw string payloads that embed JSON ---
     def safe_json_load(s):
         if s is None:
             return None
         if isinstance(s, (dict, list)):
             return s
         if isinstance(s, str):
+            s = s.strip()
+            if s == "":
+                return None
             try:
                 return json.loads(s)
             except Exception:
-                try:
-                    t = s.strip()
-                    t = t.replace("\\'", "'").replace('\"{', '{').replace('}\"', '}')
-                    return json.loads(t)
-                except Exception:
-                    return None
+                # try to extract balanced-brace substrings (simple approach)
+                subs = []
+                stack = []
+                start = None
+                for i, ch in enumerate(s):
+                    if ch == "{":
+                        if not stack:
+                            start = i
+                        stack.append("{")
+                    elif ch == "}":
+                        if stack:
+                            stack.pop()
+                            if not stack and start is not None:
+                                subs.append(s[start:i+1])
+                                start = None
+                for sub in subs:
+                    try:
+                        return json.loads(sub)
+                    except Exception:
+                        continue
         return None
 
+    # --- normalize numeric strings like "1,800.00" -> float ---
     def clean_number(x):
         if x is None:
             return 0.0
@@ -263,6 +282,7 @@ def map_prediction_to_ui(pred):
         s = str(x).strip()
         if s == "":
             return 0.0
+        # remove commas and non-number chars except dot and minus
         s = re.sub(r"[,\s]", "", s)
         s = re.sub(r"[^\d\.\-]", "", s)
         if s in ("", ".", "-", "-."):
@@ -272,28 +292,69 @@ def map_prediction_to_ui(pred):
         except Exception:
             return 0.0
 
-    def parse_date(s):
-        if not s:
-            return ""
-        s = str(s).strip()
-        for fmt in ("%Y-%m-%d", "%d-%m-%Y", "%d/%m/%Y", "%m/%d/%Y", "%d.%m.%Y"):
-            try:
-                return datetime.strptime(s, fmt).strftime("%Y-%m-%d")
-            except Exception:
-                pass
-        m = re.match(r"^(\d{1,2})/(\d{1,2})/(\d{4})$", s)
-        if m:
-            a, b, y = int(m.group(1)), int(m.group(2)), int(m.group(3))
-            if a > 12:
-                d, mo = a, b
-            else:
-                mo, d = a, b
-            try:
-                return datetime(year=y, month=mo, day=d).strftime("%Y-%m-%d")
-            except Exception:
-                return s
-        return s
+    # --- collect all keys -> list of values from pred, recursively ---
+    def collect_keys(obj, out):
+        if isinstance(obj, dict):
+            for k, v in obj.items():
+                lk = str(k).strip().lower()
+                out[lk].append(v)
+                collect_keys(v, out)
+        elif isinstance(obj, list):
+            for it in obj:
+                collect_keys(it, out)
+        else:
+            # primitive: handled via parent key
+            pass
+
+    # --- find list-of-dicts candidates for items (recursively) ---
+    def collect_lists_of_dicts(obj, out_lists):
+        if isinstance(obj, dict):
+            for v in obj.values():
+                if isinstance(v, list) and v and isinstance(v[0], dict):
+                    out_lists.append(v)
+                else:
+                    collect_lists_of_dicts(v, out_lists)
+        elif isinstance(obj, list):
+            for it in obj:
+                if isinstance(it, list) and it and isinstance(it[0], dict):
+                    out_lists.append(it)
+                else:
+                    collect_lists_of_dicts(it, out_lists)
+
+    # --- map item dict -> UI item row using the keys you specified in example ---
+    def map_item_dict(it):
+        if not isinstance(it, dict):
+            return None
+        # lowered keys mapping
+        lower = {str(k).strip().lower(): v for k, v in it.items()}
+        desc = (lower.get("descriptions") or lower.get("description") or lower.get("desc") or lower.get("item") or "")
+        qty = lower.get("quantity") or lower.get("qty") or lower.get("count") or ""
+        unit_price = lower.get("unit_price") or lower.get("price") or ""
+        amount = lower.get("amount") or lower.get("line_total") or lower.get("line total") or lower.get("total") or ""
+        tax = lower.get("tax") or lower.get("tax_amount") or ""
+        line_total = lower.get("line_total") or lower.get("line_total".lower()) or lower.get("line total") or amount
+
+        return {
+            "Description": str(desc).strip(),
+            "Quantity": float(clean_number(qty)),
+            "Unit Price": float(clean_number(unit_price)),
+            "Amount": float(clean_number(amount)),
+            "Tax": float(clean_number(tax)),
+            "Line Total": float(clean_number(line_total))
+        }
 
+    # ----------------- Start mapping -----------------
+    # Try parse if pred is a JSON-like string
+    parsed = safe_json_load(pred) if isinstance(pred, str) else pred
+    if parsed is None and isinstance(pred, str):
+        # not parseable -> fallback to empty UI
+        parsed = None
+
+    if parsed is None and not isinstance(pred, dict):
+        # nothing we can map
+        parsed = pred  # will still allow collect_keys if it's dict; else produce empty ui
+
+    # create empty UI template
     ui = {
         "Invoice Number": "",
         "Invoice Date": "",
@@ -313,137 +374,121 @@ def map_prediction_to_ui(pred):
         "Itemized Data": []
     }
 
-    if pred is None:
-        return ui
-
-    if isinstance(pred, str):
-        parsed = safe_json_load(pred)
-        if parsed is not None:
-            pred = parsed
-
-    gt = None
-    if isinstance(pred, dict):
-        if "gt_parse" in pred:
-            gp = pred["gt_parse"]
-            gp_parsed = safe_json_load(gp)
-            gt = gp_parsed if gp_parsed is not None else (gp if isinstance(gp, dict) else {})
-        else:
-            gt = pred
-    else:
-        return ui
-
-    header = gt.get("header") or {}
-    items = gt.get("items") or []
-    summary = gt.get("summary") or {}
-
-    ui["Invoice Number"] = header.get("invoice_no") or header.get("invoice_number") or ui["Invoice Number"]
-    ui["Invoice Date"] = str(header.get("invoice_date") or header.get("inv_date") or "")
-    ui["Due Date"] = str(header.get("due_date") or header.get("due") or "")
-
-    ui["Sender Name"] = header.get("sender_name") or header.get("seller_name") or header.get("from_name") or ui["Sender Name"]
-    ui["Sender Address"] = header.get("sender_addr") or header.get("sender_address") or header.get("seller_addr") or ui["Sender Address"]
-    ui["Recipient Name"] = header.get("rcpt_name") or header.get("recipient_name") or header.get("to_name") or ui["Recipient Name"]
-    ui["Recipient Address"] = header.get("rcpt_addr") or header.get("rcpt_address") or header.get("recipient_address") or ui["Recipient Address"]
+    # If parsed is a dict, collect all keys and list-of-dict candidates
+    key_map = defaultdict(list)  # lowercase-key -> list of values
+    list_candidates = []         # list of list-of-dicts found
+    if isinstance(parsed, dict):
+        collect_keys(parsed, key_map)
+        collect_lists_of_dicts(parsed, list_candidates)
+    elif isinstance(pred, dict):
+        # if parsing failed but original pred is dict, use that
+        collect_keys(pred, key_map)
+        collect_lists_of_dicts(pred, list_candidates)
+
+    # Helper to pick first non-empty value from candidate keys
+    def pick_first(*candidate_keys):
+        for k in candidate_keys:
+            lk = k.strip().lower()
+            if lk in key_map:
+                # pick first non-empty
+                for v in key_map[lk]:
+                    if v is None:
+                        continue
+                    # return primitive or string immediately; if dict/list, return as-is
+                    if isinstance(v, (dict, list)):
+                        return v
+                    s = str(v).strip()
+                    if s != "":
+                        return s
+        return None
 
-    ui["Sender"] = {"Name": ui["Sender Name"], "Address": ui["Sender Address"]}
-    ui["Recipient"] = {"Name": ui["Recipient Name"], "Address": ui["Recipient Address"]}
+    # Map simple scalar fields using the exact keys you provided (plus common close variants)
+    ui["Invoice Number"] = pick_first("invoice_no", "invoice_number", "invoiceid", "invoice id") or ""
+    ui["Invoice Date"] = pick_first("invoice_date", "date", "invoice date") or ""
+    ui["Due Date"] = pick_first("due_date", "due_date", "due") or ""
+    ui["Sender Name"] = pick_first("sender_name", "sender") or ""
+    ui["Sender Address"] = pick_first("sender_addr", "sender_address", "sender addr") or ""
+    ui["Recipient Name"] = pick_first("rcpt_name", "recipient_name", "recipient", "rcpt") or ""
+    ui["Recipient Address"] = pick_first("rcpt_addr", "recipient_address", "recipient addr") or ""
 
+    # bank details: gather keys that start with 'bank_' or exact matches
     bank = {}
-    if header.get("bank_name"):
-        bank["bank_name"] = str(header.get("bank_name")).strip()
-    if header.get("bank_acc_no"):
-        bank["bank_account_number"] = str(header.get("bank_acc_no")).strip()
-    if header.get("bank_account_number"):
-        bank["bank_account_number"] = bank.get("bank_account_number") or str(header.get("bank_account_number")).strip()
-    if header.get("bank_iban"):
-        bank["bank_iban"] = str(header.get("bank_iban")).strip()
-    if header.get("bank_routing"):
-        bank["bank_routing"] = str(header.get("bank_routing")).strip()
-    if header.get("bank_swift"):
-        bank["bank_swift"] = str(header.get("bank_swift")).strip()
-    if header.get("bank_branch"):
-        bank["bank_branch"] = str(header.get("bank_branch")).strip()
-    if header.get("bank_acc_name"):
-        bank["bank_acc_name"] = str(header.get("bank_acc_name")).strip()
-    hb = header.get("bank")
-    if isinstance(hb, dict):
-        for k, v in hb.items():
-            if not v:
-                continue
-            lk = k.lower()
-            if "iban" in lk:
-                bank["bank_iban"] = bank.get("bank_iban") or str(v).strip()
-            elif "swift" in lk:
-                bank["bank_swift"] = bank.get("bank_swift") or str(v).strip()
-            elif "acc" in lk or "account" in lk:
-                bank["bank_account_number"] = bank.get("bank_account_number") or str(v).strip()
-            elif "name" in lk and "bank" in lk:
-                bank["bank_name"] = bank.get("bank_name") or str(v).strip()
-            elif "branch" in lk:
-                bank["bank_branch"] = bank.get("bank_branch") or str(v).strip()
-            elif "acc_name" in lk or "account_name" in lk:
-                bank["bank_acc_name"] = bank.get("bank_acc_name") or str(v).strip()
-
+    for bk in ("bank_name", "bank_acc_no", "bank_account_number", "bank_acc_name", "bank_iban", "bank_swift", "bank_routing", "bank_branch", "iban"):
+        val = pick_first(bk, bk.replace("bank_", ""))  # allow both 'iban' and 'bank_iban'
+        if val:
+            # normalize key name to bank_* form
+            if bk == "iban":
+                bank["bank_iban"] = str(val)
+            else:
+                bank[bk] = str(val)
     ui["Bank Details"] = bank
 
-    ui["Subtotal"] = clean_number(summary.get("subtotal") or summary.get("sub_total") or summary.get("subTotal"))
-    ui["Tax Percentage"] = clean_number(summary.get("tax_rate") or summary.get("taxRate") or summary.get("tax_percentage"))
-    ui["Total Tax"] = clean_number(summary.get("tax_amount") or summary.get("tax") or summary.get("taxAmount"))
-    ui["Total Amount"] = clean_number(summary.get("total_amount") or summary.get("grand_total") or summary.get("total") or summary.get("amount_total"))
-    ui["Currency"] = summary.get("currency") or header.get("currency") or ui["Currency"] or ""
-
-    normalized_items = []
-
-    if isinstance(items, str):
-        parsed_items = safe_json_load(items)
-        if parsed_items is not None:
-            items = parsed_items
-
-    if isinstance(items, dict):
-        if any(isinstance(v, list) for v in items.values()):
-            list_cols = {k: v for k, v in items.items() if isinstance(v, list)}
-            max_len = max((len(v) for v in list_cols.values()), default=0)
-            for i in range(max_len):
-                row = {}
-                for k, v in items.items():
-                    if isinstance(v, list):
-                        row[k] = v[i] if i < len(v) else ""
-                    else:
-                        row[k] = v
-                normalized_items.append(row)
-        else:
-            normalized_items.append(items)
-    elif isinstance(items, list):
-        normalized_items = items
-    else:
-        normalized_items = []
-
-    item_rows = []
-    for it in normalized_items:
-        if not isinstance(it, dict):
-            item_rows.append({"Description": str(it), "Quantity": 1, "Unit Price": 0.0, "Amount": 0.0, "Tax": 0.0, "Line Total": 0.0})
-            continue
-        desc = it.get("descriptions") or it.get("description") or it.get("desc") or it.get("item") or it.get("name") or ""
-        qty = it.get("quantity") or it.get("qty") or it.get("Quantity") or ""
-        unit = it.get("unit_price") or it.get("unitPrice") or it.get("price") or ""
-        amt = it.get("amount") or it.get("Line_total") or it.get("line_total") or it.get("total") or ""
-
-        # Extract item-level tax if available under common keys
-        tax_val = it.get("tax") or it.get("tax_amount") or it.get("line_tax") or it.get("item_tax") or it.get("taxAmount") or ""
-
-        # Extract explicit line total if present; otherwise fall back to amount
-        line_total_val = it.get("Line_total") or it.get("line_total") or it.get("lineTotal") or amt
-
-        item_rows.append({
-            "Description": str(desc).strip(),
-            "Quantity": float(clean_number(qty)),
-            "Unit Price": float(clean_number(unit)),
-            "Amount": float(clean_number(amt)),
-            "Tax": float(clean_number(tax_val)),
-            "Line Total": float(clean_number(line_total_val))
-        })
-
-    ui["Itemized Data"] = item_rows
+    # summary / totals
+    ui["Subtotal"] = clean_number(pick_first("subtotal", "sub_total", "sub total") or 0.0)
+    ui["Tax Percentage"] = clean_number(pick_first("tax_rate", "tax_percentage", "tax pct", "tax percentage") or 0.0)
+    ui["Total Tax"] = clean_number(pick_first("tax_amount", "tax", "total_tax") or 0.0)
+    ui["Total Amount"] = clean_number(pick_first("total_amount", "grand_total", "total", "amount") or 0.0)
+    ui["Currency"] = (pick_first("currency") or "").strip()
+
+    # Item extraction:
+    items_rows = []
+
+    # 1) If we found an explicit list-of-dicts candidate, use the first one that looks like invoice items
+    def list_looks_like_items(lst):
+        if not isinstance(lst, list) or not lst:
+            return False
+        if not isinstance(lst[0], dict):
+            return False
+        # check if any expected item key present in first element
+        expected = {"descriptions", "description", "desc", "item", "quantity", "qty", "amount", "unit_price", "line_total", "line_total".lower(), "line_total"}
+        keys0 = {str(k).strip().lower() for k in lst[0].keys()}
+        return bool(expected.intersection(keys0))
+
+    for cand in list_candidates:
+        if list_looks_like_items(cand):
+            for it in cand:
+                row = map_item_dict(it)
+                if row is not None:
+                    items_rows.append(row)
+            # prefer first plausible list
+            if items_rows:
+                break
+
+    # 2) If no list-of-dicts found, try to find a single dict anywhere that looks like an item (e.g., 'items': {...} as dict)
+    if not items_rows:
+        # search key_map values for dicts that have item-like keys
+        for k, vals in key_map.items():
+            for v in vals:
+                if isinstance(v, dict):
+                    # does this dict have an item-like key?
+                    lower_keys = {str(x).strip().lower() for x in v.keys()}
+                    if lower_keys.intersection({"descriptions", "description", "desc", "amount", "line_total", "quantity", "qty", "unit_price"}):
+                        row = map_item_dict(v)
+                        if row is not None:
+                            items_rows.append(row)
+                            # we don't break because there might be multiple item-like dicts at different keys,
+                            # but continue scanning to collect all.
+    # 3) Last resort: if key_map contains 'descriptions' or 'amount' as scalar but no dict, build a single-item row
+    if not items_rows:
+        desc = pick_first("descriptions", "description")
+        amt = pick_first("amount", "line_total")
+        qty = pick_first("quantity", "qty")
+        unit_price = pick_first("unit_price", "price")
+        if desc or amt or qty or unit_price:
+            items_rows.append({
+                "Description": str(desc or ""),
+                "Quantity": float(clean_number(qty)),
+                "Unit Price": float(clean_number(unit_price)),
+                "Amount": float(clean_number(amt)),
+                "Tax": float(clean_number(pick_first("tax", "tax_amount") or 0.0)),
+                "Line Total": float(clean_number(amt or 0.0))
+            })
+
+    ui["Itemized Data"] = items_rows
+
+    # Also set Sender/Recipient convenience fields
+    ui["Sender"] = {"Name": ui["Sender Name"], "Address": ui["Sender Address"]}
+    ui["Recipient"] = {"Name": ui["Recipient Name"], "Address": ui["Recipient Address"]}
 
     return ui