Update src/streamlit_app.py

src/streamlit_app.py

@@ -196,7 +196,6 @@ def run_inference_on_image(image: Image.Image, processor, model, device, decoder
 
 # ---------------------------
 # Helper: map donut output to our UI schema
-# (kept unchanged from your original)
 # ---------------------------
 def map_prediction_to_ui(pred):
     import json, re
@@ -218,77 +217,6 @@ def map_prediction_to_ui(pred):
                 except Exception:
                     return None
         return None
-def flatten_invoice_to_rows(invoice_data) -> list:
-    """
-    Converts nested invoice data into a flat list of rows (one per line item),
-    with invoice-level and sender/recipient/bank fields repeated in each row.
-    """
-    rows = []
-    line_items = invoice_data.get("Itemized Data", [])
-    if not line_items:
-        # If no line items, create one row with invoice info only
-        row = {
-            "Invoice Number": invoice_data.get("Invoice Number", ""),
-            "Invoice Date": invoice_data.get("Invoice Date", ""),
-            "Due Date": invoice_data.get("Due Date", ""),
-            "Currency": invoice_data.get("Currency", ""),
-            "Subtotal": invoice_data.get("Subtotal", 0.0),
-            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
-            "Total Tax": invoice_data.get("Total Tax", 0.0),
-            "Total Amount": invoice_data.get("Total Amount", 0.0),
-            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
-            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
-            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
-            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
-        }
-        # Flatten bank details
-        bank = invoice_data.get("Bank Details", {})
-        for k, v in bank.items():
-            row[f"bank_{k}"] = v
-
-        # Add empty line item fields
-        row.update({
-            "Item Description": "",
-            "Item Quantity": 0,
-            "Item Unit Price": 0.0,
-            "Item Amount": 0.0,
-        })
-        rows.append(row)
-        return rows
-
-    # For each line item, create a row with all invoice context
-    for item in line_items:
-        row = {
-            "Invoice Number": invoice_data.get("Invoice Number", ""),
-            "Invoice Date": invoice_data.get("Invoice Date", ""),
-            "Due Date": invoice_data.get("Due Date", ""),
-            "Currency": invoice_data.get("Currency", ""),
-            "Subtotal": invoice_data.get("Subtotal", 0.0),
-            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
-            "Total Tax": invoice_data.get("Total Tax", 0.0),
-            "Total Amount": invoice_data.get("Total Amount", 0.0),
-            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
-            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
-            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
-            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
-        }
-
-        # Flatten bank details
-        bank = invoice_data.get("Bank Details", {})
-        for k, v in bank.items():
-            row[f"bank_{k}"] = v
-
-        # Add line item fields
-        row.update({
-            "Item Description": item.get("Description", ""),
-            "Item Quantity": item.get("Quantity", 0),
-            "Item Unit Price": item.get("Unit Price", 0.0),
-            "Item Amount": item.get("Amount", 0.0),
-        })
-
-        rows.append(row)
-
-    return rows
 
     def clean_number(x):
         if x is None:
@@ -474,6 +402,81 @@ def flatten_invoice_to_rows(invoice_data) -> list:
 
     return ui
 
+# ---------------------------
+# Helper: flatten invoice to CSV rows
+# ---------------------------
+def flatten_invoice_to_rows(invoice_data) -> list:
+    """
+    Converts nested invoice data into a flat list of rows (one per line item),
+    with invoice-level and sender/recipient/bank fields repeated in each row.
+    """
+    rows = []
+    line_items = invoice_data.get("Itemized Data", [])
+    if not line_items:
+        # If no line items, create one row with invoice info only
+        row = {
+            "Invoice Number": invoice_data.get("Invoice Number", ""),
+            "Invoice Date": invoice_data.get("Invoice Date", ""),
+            "Due Date": invoice_data.get("Due Date", ""),
+            "Currency": invoice_data.get("Currency", ""),
+            "Subtotal": invoice_data.get("Subtotal", 0.0),
+            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
+            "Total Tax": invoice_data.get("Total Tax", 0.0),
+            "Total Amount": invoice_data.get("Total Amount", 0.0),
+            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
+            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
+            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
+            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
+        }
+        # Flatten bank details
+        bank = invoice_data.get("Bank Details", {})
+        for k, v in bank.items():
+            row[f"bank_{k}"] = v
+
+        # Add empty line item fields
+        row.update({
+            "Item Description": "",
+            "Item Quantity": 0,
+            "Item Unit Price": 0.0,
+            "Item Amount": 0.0,
+        })
+        rows.append(row)
+        return rows
+
+    # For each line item, create a row with all invoice context
+    for item in line_items:
+        row = {
+            "Invoice Number": invoice_data.get("Invoice Number", ""),
+            "Invoice Date": invoice_data.get("Invoice Date", ""),
+            "Due Date": invoice_data.get("Due Date", ""),
+            "Currency": invoice_data.get("Currency", ""),
+            "Subtotal": invoice_data.get("Subtotal", 0.0),
+            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
+            "Total Tax": invoice_data.get("Total Tax", 0.0),
+            "Total Amount": invoice_data.get("Total Amount", 0.0),
+            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
+            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
+            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
+            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
+        }
+
+        # Flatten bank details
+        bank = invoice_data.get("Bank Details", {})
+        for k, v in bank.items():
+            row[f"bank_{k}"] = v
+
+        # Add line item fields
+        row.update({
+            "Item Description": item.get("Description", ""),
+            "Item Quantity": item.get("Quantity", 0),
+            "Item Unit Price": item.get("Unit Price", 0.0),
+            "Item Amount": item.get("Amount", 0.0),
+        })
+
+        rows.append(row)
+
+    return rows
+
 # Load model once
 try:
     with st.spinner("Loading model & processor (cached) ..."):
@@ -553,7 +556,11 @@ if not st.session_state.is_processing_batch and len(st.session_state.batch_resul
                 mapped = map_prediction_to_ui(pred)
             except Exception as e:
                 st.warning(f"Error processing {uploaded_file.name}: {str(e)}")
-                pred, mapped = None, {}
+                pred = None
+                mapped = {}  # 👈 Ensure mapped is always a dict
+
+            # ✅ SAFETY: Ensure mapped is a dict before copying
+            safe_mapped = mapped if isinstance(mapped, dict) else {}
 
             # Save to session state
             st.session_state.batch_results[file_hash] = {
@@ -561,7 +568,7 @@ if not st.session_state.is_processing_batch and len(st.session_state.batch_resul
                 "image": image,
                 "raw_pred": pred,
                 "mapped_data": mapped,
-                "edited_data": mapped.copy()  # editable copy
+                "edited_data": safe_mapped.copy()  # editable copy — now safe
             }
 
             progress_bar.progress((idx + 1) / len(uploaded_files))
@@ -711,6 +718,9 @@ elif len(st.session_state.batch_results) > 0:
             data['Bank Details'] = bank_info
 
         # ---------- Line Items ----------
+                # ---------- Line Items ----------
+                # ---------- Line Items ----------
+                # ---------- Line Items ----------
         with tabs[3]:
             editor_key = f"item_editor_{selected_hash}"
 
@@ -720,29 +730,27 @@ elif len(st.session_state.batch_results) > 0:
                 if col not in df.columns:
                     df[col] = ""
 
-            # ✅ FIXED: Use on_change to force immediate rerun → edited_df is always fresh
+            st.write("✏️ Edit line items below. Press Enter or click outside a cell to confirm each edit.")
+            # Get the edited DataFrame directly from data_editor
             edited_df = st.data_editor(
                 df,
                 num_rows="dynamic",
                 key=editor_key,
                 use_container_width=True,
-                on_change=trigger_rerun  # 👈 This is the magic fix
+                on_change=trigger_rerun
             )
 
-            # ✅ Safe to use — edited_df is updated after forced rerun
-            data['Itemized Data'] = edited_df.to_dict('records')
-
-            # ❗ OPTIONAL: Auto-calculate Amount = Quantity × Unit Price
-            # Uncomment below if you want auto-calculation:
-            # if "Quantity" in edited_df.columns and "Unit Price" in edited_df.columns:
-            #     edited_df["Amount"] = (edited_df["Quantity"] * edited_df["Unit Price"]).round(2)
-            #     data['Itemized Data'] = edited_df.to_dict('records')
-
             if len(edited_df) == 0:
                 st.info("No line items found in the invoice.")
 
-        # Save button (per file) — OPTIONAL, since edits auto-save via reference
+        # ... (previous code) ...
+
+        # Save button (per file)
         if st.button("💾 Save Edits for This File", key=f"save_{selected_hash}"):
+            # Update line items from the current edited_df (which is a DataFrame)
+            data['Itemized Data'] = edited_df.to_dict('records')
+    
+            # Save the entire data to session state
             st.session_state.batch_results[selected_hash]["edited_data"] = data
             st.success(f"✅ Edits saved for {current['file_name']}")
 
@@ -798,7 +806,7 @@ elif len(st.session_state.batch_results) > 0:
                 json_name = f"{Path(result['file_name']).stem}_extracted.json"
                 zf.writestr(json_name, json_data)
 
-                                # Save FULL CSV (all data)
+                # Save FULL CSV (all data)
                 rows = flatten_invoice_to_rows(result["edited_data"])
                 full_df = pd.DataFrame(rows)