Update app.py

app.py

@@ -1,121 +1,147 @@
-import streamlit as st
-import pandas as pd
-from azure.core.credentials import AzureKeyCredential
-from azure.ai.formrecognizer import DocumentAnalysisClient
-import time
-from dotenv import load_dotenv
-import os
-
-load_dotenv()
-
-# Azure Form Recognizer credentials
-endpoint = st.secrets["endpoint"]
-key = st.secrets["key"]
-
-# Initialize DocumentAnalysisClient
-document_analysis_client = DocumentAnalysisClient(
-    endpoint=endpoint, credential=AzureKeyCredential(key)
-)
-
-# Sidebar for document type selection
-document_type = st.sidebar.selectbox(
-    "Select Document Type", ["Invoice", "Receipt", "Identity Document"]
-)
-
-# Center area for file upload
-uploaded_file = st.file_uploader("Choose or drag a file", type=["pdf", "png", "jpg", "jpeg"])
-
-# Initialize session state for processing results
-if "processed" not in st.session_state:
-    st.session_state.processed = False
-    st.session_state.tables = []
-    st.session_state.kv_df = pd.DataFrame()
-
-# Process the uploaded file if the button is clicked
-if uploaded_file is not None:
-    if st.button("Upload and Process"):
-        with st.spinner("Processing..."):
-            time.sleep(1)  # Simulate some delay to show the spinner
-
-            # Analyze the document
-            poller = document_analysis_client.begin_analyze_document(
-                "prebuilt-document", document=uploaded_file
-            )
-            result = poller.result()
-
-            # Extract key-value pairs and store them in a dictionary
-            kv_dict = {}
-            for kv_pair in result.key_value_pairs:
-                if kv_pair.key and kv_pair.value:  # Ensure both key and value exist
-                    kv_dict[kv_pair.key.content] = kv_pair.value.content
-
-            kv_df = pd.DataFrame(list(kv_dict.items()), columns=["Key", "Value"]).T
-
-            # Set the first row as the header
-            header = kv_df.iloc[0]
-
-            # Create a new DataFrame with the header and the remaining rows
-            kv_df = kv_df[1:].reset_index(drop=True)
-
-            # Insert the new header row
-            header_df = pd.DataFrame([header], columns=range(len(header)))
-
-            # Concatenate the header row DataFrame with the original DataFrame
-            kv_df = pd.concat([header_df, kv_df], ignore_index=True)
-
-            st.session_state.kv_df = kv_df
-            st.session_state.tables = []
-
-            if result.tables:
-                for table in result.tables:
-                    temp_kv_df = kv_df.copy()
-
-                    data = []
-                    for cell in table.cells:
-                        data.append([cell.row_index, cell.column_index, cell.content])
-
-                    table_df = pd.DataFrame(data, columns=["row_index", "column_index", "content"])
-                    table_df = table_df.pivot(index="row_index", columns="column_index", values="content")
-
-                    rows_to_add = len(table_df) - len(temp_kv_df)
-                    if rows_to_add > 0:
-                        last_row = temp_kv_df.iloc[-1]
-                        additional_rows = pd.DataFrame([last_row] * rows_to_add, columns=temp_kv_df.columns)
-                        temp_kv_df = pd.concat([temp_kv_df, additional_rows], ignore_index=True)
-
-                    table_with_kv = pd.concat([temp_kv_df, table_df], axis=1, ignore_index=True)
-                    st.session_state.tables.append(table_with_kv)
-
-            st.session_state.processed = True
-
-# Display the results if processing is done
-if st.session_state.processed:
-    st.write("Extracted Key-Value Pairs:")
-    st.dataframe(st.session_state.kv_df)
-
-    if st.session_state.tables:
-        for i, table_with_kv in enumerate(st.session_state.tables):
-            st.write(f"Table {i + 1} with Key-Value Pairs:")
-            st.dataframe(table_with_kv)
-
-            st.download_button(
-                label=f"Download Table {i + 1} as CSV",
-                data=table_with_kv.to_csv(index=False, header=False).encode('utf-8'),
-                file_name=f"table_with_kv_{i + 1}.csv",
-                mime='text/csv',
-            )
-    else:
-        st.write("No tables found in the document.")
-        st.download_button(
-            label="Download Key-Value Pairs as CSV",
-            data=st.session_state.kv_df.to_csv(index=False, header=False).encode('utf-8'),
-            file_name="kv_df.csv",
-            mime='text/csv',
-        )
-
-# Option to remove the file and clear the session
-if st.session_state.processed and st.button("Remove File"):
-    st.session_state.processed = False
-    st.session_state.tables = []
-    st.session_state.kv_df = pd.DataFrame()
-    st.rerun()
+import streamlit as st
+import pandas as pd
+from azure.core.credentials import AzureKeyCredential
+from azure.ai.formrecognizer import DocumentAnalysisClient
+import time
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+# Azure Form Recognizer credentials
+endpoint = st.secrets["endpoint"]
+key = st.secrets["key"]
+
+# Initialize DocumentAnalysisClient
+document_analysis_client = DocumentAnalysisClient(
+    endpoint=endpoint, credential=AzureKeyCredential(key)
+)
+
+# Sidebar for document type selection
+document_type = st.sidebar.selectbox(
+    "Select Document Type", ["Invoice", "Receipt", "Identity Document"]
+)
+
+# Center area for file upload
+uploaded_file = st.file_uploader(f"Choose or drag an {document_type} file", type=["pdf", "png", "jpg", "jpeg"])
+
+# Initialize session state for processing results
+if "processed" not in st.session_state:
+    st.session_state.processed = False
+    st.session_state.tables = []
+    st.session_state.kv_df = pd.DataFrame()
+
+# Process the uploaded file if the button is clicked
+if uploaded_file is not None:
+    if st.button("Upload and Process"):
+        with st.spinner("Processing..."):
+            time.sleep(1)  # Simulate some delay to show the spinner
+
+            # Analyze the document
+            poller = document_analysis_client.begin_analyze_document(
+                "prebuilt-document", document=uploaded_file
+            )
+            result = poller.result()
+
+            # Extract key-value pairs and store them in a dictionary
+            kv_dict = {}
+            for kv_pair in result.key_value_pairs:
+                if kv_pair.key and kv_pair.value:  # Ensure both key and value exist
+                    kv_dict[kv_pair.key.content] = kv_pair.value.content
+
+            kv_df = pd.DataFrame(list(kv_dict.items()), columns=["Key", "Value"]).T
+
+            # Set the first row as the header
+            header = kv_df.iloc[0]
+
+            # Create a new DataFrame with the header and the remaining rows
+            kv_df = kv_df[1:].reset_index(drop=True)
+
+            # Insert the new header row
+            header_df = pd.DataFrame([header], columns=range(len(header)))
+
+            # Concatenate the header row DataFrame with the original DataFrame
+            kv_df = pd.concat([header_df, kv_df], ignore_index=True)
+
+            st.session_state.kv_df = kv_df
+            st.session_state.tables = []
+
+            if result.tables:
+                for table in result.tables:
+                    temp_kv_df = kv_df.copy()
+
+                    data = []
+                    for cell in table.cells:
+                        data.append([cell.row_index, cell.column_index, cell.content])
+
+                    table_df = pd.DataFrame(data, columns=["row_index", "column_index", "content"])
+                    table_df = table_df.pivot(index="row_index", columns="column_index", values="content")
+
+                    st.session_state.tables.append(table_df)
+
+            st.session_state.processed = True
+
+# Display the results if processing is done
+if st.session_state.processed:
+    if not st.session_state.kv_df.empty:
+        st.write("Extracted Key-Value Pairs:")
+        st.dataframe(st.session_state.kv_df)
+        st.download_button(
+            label="Download Key-Value Pairs as CSV",
+            data=st.session_state.kv_df.to_csv(index=False, header=False).encode('utf-8'),
+            file_name="kv_df.csv",
+            mime='text/csv',
+        )
+    else:
+        st.write("No Key-Value Pairs found in document.")
+
+    
+    if st.session_state.tables:
+        if not st.session_state.kv_df.empty:
+            # Allow the user to select multiple columns
+            first_row_values = st.session_state.kv_df.iloc[0].values
+            selected_columns = st.multiselect(
+                "Select columns from Key-Value Pairs to add to tables:", first_row_values
+            )
+        for i, table_df in enumerate(st.session_state.tables):
+            st.write(f"Table {i + 1} with Selected Columns Added:")
+
+            # Find the column names corresponding to the selected first row values
+            column_indices = [list(first_row_values).index(val) for val in selected_columns]
+            column_names = st.session_state.kv_df.columns[column_indices]
+            
+            # Add the selected columns to the table
+            if selected_columns:
+                # Create a new DataFrame for the selected columns
+                selected_columns_df = st.session_state.kv_df[column_names].copy()
+                
+                # Align the length of selected_columns_df with the table_df rows
+                selected_columns_df = selected_columns_df.iloc[:len(table_df)].reset_index(drop=True)
+                
+                rows_to_add = len(table_df) - len(selected_columns_df)
+                if rows_to_add > 0:
+                    last_row = selected_columns_df.iloc[-1]
+                    additional_rows = pd.DataFrame([last_row] * rows_to_add, columns=selected_columns_df.columns)
+                    selected_columns_df = pd.concat([selected_columns_df, additional_rows], ignore_index=True)
+                    
+                # Concatenate the selected columns DataFrame with the table DataFrame
+                table_df = pd.concat([selected_columns_df, table_df.reset_index(drop=True)], axis=1, ignore_index=True)
+
+            st.dataframe(table_df)
+
+            st.download_button(
+                label=f"Download Table {i + 1} as CSV",
+                data=table_df.to_csv(index=False, header=False).encode('utf-8'),
+                file_name=f"table_with_kv_{i + 1}.csv",
+                mime='text/csv',
+            )
+    else:
+        st.write("No tables found in the document.")
+        
+
+# Option to remove the file and clear the session
+if st.session_state.processed and st.button("Remove File"):
+    st.session_state.processed = False
+    st.session_state.tables = []
+    st.session_state.kv_df = pd.DataFrame()
+    st.rerun()