Spaces:
Build error
Build error
| 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 | |
| from dotenv import load_dotenv | |
| import os | |
| load_dotenv() | |
| # Azure Form Recognizer credentials | |
| endpoint = os.environ["endpoint"] | |
| key = os.environ["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 | |
| clean_key = kv_pair.key.content.rstrip(':') | |
| kv_dict[clean_key] = 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: | |
| 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") | |
| # 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.write("Key-Value Pairs found.") | |
| else: | |
| st.sidebar.write("No Key-Value Pairs found in document.") | |
| if st.session_state.tables: | |
| # Display available tables for selection in the sidebar | |
| table_options = ["Select"] + [f"Table {i + 1}" for i in range(len(st.session_state.tables))] | |
| selected_table_option = st.sidebar.selectbox( | |
| "Select table to join key-value pairs:", table_options, index=0 | |
| ) | |
| # Proceed only if a table is selected | |
| if selected_table_option != "Select": | |
| 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: | |
| if col not in selected_table_df.columns: | |
| selected_table_df[col] = np.nan | |
| # Move default columns to the left | |
| # We include only those columns that exist in the table | |
| existing_columns = [col for col in default_columns if col in selected_table_df.columns] | |
| remaining_columns = [col for col in selected_table_df.columns if col not in existing_columns] | |
| selected_table_df = selected_table_df[[*existing_columns, *remaining_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 | |
| second_row_values = st.session_state.kv_df.iloc[1].values | |
| key_value_options = [ | |
| f"{key}: {value}" for key, value in zip(first_row_values, second_row_values) | |
| ] | |
| selected_columns = [] | |
| st.sidebar.write("Select key-value pairs to the default columns:") | |
| # 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)) | |
| # 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] | |
| st.write(f"Selected Table {selected_table_index + 1} with Key-Value Pairs Added:") | |
| # st.dataframe(selected_table_df) | |
| # Allow row and column selection for deletion | |
| st.sidebar.subheader("Manage Table") | |
| # Delete columns | |
| if st.sidebar.checkbox("Delete Columns"): | |
| col = list(selected_table_df.columns) | |
| col = col[3:] | |
| cols_to_delete = st.sidebar.multiselect( | |
| "Select Columns to Delete", col | |
| ) | |
| if st.sidebar.button("Remove Columns"): | |
| st.session_state.tables[selected_table_index].drop(columns=cols_to_delete, inplace=True) | |
| # Delete rows | |
| if st.sidebar.checkbox("Delete Rows"): | |
| row_indices = st.sidebar.multiselect( | |
| "Select Rows to Delete", selected_table_df.index.tolist() | |
| ) | |
| if st.sidebar.button("Remove Rows"): | |
| st.session_state.tables[selected_table_index].drop(index=row_indices, inplace=True) | |
| # Show the updated table | |
| # st.write(f"Updated Table {selected_table_index + 1}:") | |
| st.dataframe(selected_table_df) | |
| # Download updated table | |
| st.download_button( | |
| label=f"Download Updated Table {selected_table_index + 1} as CSV", | |
| data=selected_table_df.to_csv(index=False).encode('utf-8'), | |
| file_name=f"updated_table_{selected_table_index + 1}.csv", | |
| mime='text/csv', | |
| ) | |
| else: | |
| st.sidebar.write("No tables found in the document.") | |
| # Option to remove the file and clear the session | |
| if st.session_state.processed and st.sidebar.button("Remove File"): | |
| st.session_state.processed = False | |
| st.session_state.tables = [] | |
| st.session_state.kv_df = pd.DataFrame() | |
| st.rerun() | |