Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import pandas as pd
+import numpy as np
 from azure.core.credentials import AzureKeyCredential
 from azure.ai.formrecognizer import DocumentAnalysisClient
 import time
@@ -74,20 +75,33 @@ if uploaded_file is not None:
 
                     table_df = pd.DataFrame(data, columns=["row_index", "column_index", "content"])
                     table_df = table_df.pivot(index="row_index", columns="column_index", values="content")
-                    
+                    # Reset the column index
+                    table_df.reset_index(drop=True, inplace=True)
+
+                    # Set the first row as header
+                    table_df.columns = table_df.iloc[0]  # Take the first row as column names
+                    table_df = table_df.drop(table_df.index[0])  # Remove the first row after setting it as headers
+
+                    # Optionally reset the index if required
+                    table_df.reset_index(drop=True, inplace=True)
+                                        
                     st.session_state.tables.append(table_df)
 
             st.session_state.processed = True
 
+# Predefine default columns in the selected table
+default_columns = ['Invoice No', 'Invoice Date', 'Customer Name']
+
 # Sidebar for table and key-value pair selection
 if st.session_state.processed:
     if not st.session_state.kv_df.empty:
-        st.sidebar.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',
-        )
+        # st.sidebar.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',
+        # )
+        pass
     else:
         st.sidebar.write("No Key-Value Pairs found in document.")
 
@@ -104,6 +118,13 @@ if st.session_state.processed:
             selected_table_index = table_options.index(selected_table_option) - 1  # Adjust for "None"
             selected_table_df = st.session_state.tables[selected_table_index]
 
+            # Add default columns to the left side of the table with NaN values initially
+            for col in default_columns:
+                selected_table_df[col] = np.nan
+            
+            # Move default columns to the left
+            selected_table_df = selected_table_df[[*default_columns, *selected_table_df.columns[:-len(default_columns)]]]
+
             if not st.session_state.kv_df.empty:
                 # Allow the user to select multiple columns using checkboxes
                 first_row_values = st.session_state.kv_df.iloc[0].values
@@ -114,37 +135,32 @@ if st.session_state.processed:
                 
                 selected_columns = []
                 
-                st.sidebar.write("Select key-value pairs to add to the selected table:")
-                
-                # Create checkboxes for each key-value pair
-                for i, option in enumerate(key_value_options):
-                    if st.sidebar.checkbox(option):
-                        selected_columns.append(i)
+                st.sidebar.write("Select key-value pairs to the default columns:")
 
-                st.write(f"Selected Table {selected_table_index + 1} with Key-Value Pairs Added:")
-
-                # Add the selected columns to the table
-                if selected_columns:
-                    # Create a DataFrame for the selected columns using the original kv_df
-                    selected_columns_df = st.session_state.kv_df.iloc[:, selected_columns].copy()
-
-                    # Align the length of selected_columns_df with the table_df rows
-                    selected_columns_df = selected_columns_df.iloc[:len(selected_table_df)].reset_index(drop=True)
+                # Create draggable interface for each key-value pair
+                for i, option in enumerate(key_value_options):
+                    selected_column_option = st.sidebar.selectbox(
+                        f"Insert {option} to the column:",
+                        ["Select"] + default_columns,
+                        key=f"col_select_{i}"
+                    )
+                    
+                    if selected_column_option != "Select":
+                        selected_columns.append((i, selected_column_option))
 
-                    # Add more rows if needed to match the table_df size
-                    rows_to_add = len(selected_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)
+                # Insert selected key-value pairs into the appropriate columns
+                for idx, selected_column in selected_columns:
+                    st.session_state.kv_df[selected_column] = second_row_values[idx]
+                    
+                    # Fill the column in the table with the selected key-value pair
+                    selected_table_df[selected_column] = st.session_state.kv_df.iloc[1, idx]
 
-                    # Concatenate the selected columns with the selected table DataFrame
-                    selected_table_df = pd.concat([selected_columns_df, selected_table_df.reset_index(drop=True)], axis=1, ignore_index=True)
+                st.write(f"Selected Table {selected_table_index + 1} with Key-Value Pairs Added:")
 
                 st.dataframe(selected_table_df)
                 st.download_button(
                     label=f"Download Table {selected_table_index + 1} as CSV",
-                    data=selected_table_df.to_csv(index=False, header=False).encode('utf-8'),
+                    data=selected_table_df.to_csv(index=False).encode('utf-8'),
                     file_name=f"table_with_kv_{selected_table_index + 1}.csv",
                     mime='text/csv',
                 )
@@ -152,10 +168,10 @@ if st.session_state.processed:
                 st.write(f"Selected Table {selected_table_index + 1}:")
                 st.dataframe(selected_table_df)
                 st.download_button(
-                label=f"Download Table {selected_table_index + 1} as CSV",
-                data=selected_table_df.to_csv(index=False, header=False).encode('utf-8'),
-                file_name=f"table{selected_table_index + 1}.csv",
-                mime='text/csv',
+                    label=f"Download Table {selected_table_index + 1} as CSV",
+                    data=selected_table_df.to_csv(index=False).encode('utf-8'),
+                    file_name=f"table{selected_table_index + 1}.csv",
+                    mime='text/csv',
                 ) 
     else:
         st.sidebar.write("No tables found in the document.")