Update app.py

app.py

@@ -1,275 +1,280 @@
-import streamlit as st
-from upload_file_to_s3 import upload_file
-import base64
-import httpx
-from extract_table_from_image import process_image_using_llm
-from process_pdf import process_pdf
-from pymongo import MongoClient
-from datetime import datetime
-from table_analysis_for_image import view_table_analysis_page
-from table_analysis_for_pdf import view_pdf_table_analysis_page
-from table_analysis_for_excel import display_csv_analysis
-from view_excel import view_excel
-from copy import deepcopy
-
-import uuid
-import os
-import csv
-from view_pdf import view_pdfs
-from view_image import view_images
-from io import StringIO, BytesIO
-from dotenv import load_dotenv
-import boto3
-import pandas as pd
-st.set_page_config(layout='wide',page_title="MoSPI", page_icon="📄")
-load_dotenv()
-
-AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
-AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
-AWS_BUCKET_NAME = os.getenv("AWS_BUCKET_NAME")
-MONGO_URI = os.getenv("MONGO_URI")
-DB_NAME = os.getenv("DB_NAME")
-COLLECTION_NAME = os.getenv("COLLECTION_NAME")
-
-mongo_client = MongoClient(MONGO_URI)
-db = mongo_client[DB_NAME]
-collection = db[COLLECTION_NAME]
-
-s3 = boto3.client(
-    's3',
-    aws_access_key_id=AWS_ACCESS_KEY_ID,
-    aws_secret_access_key=AWS_SECRET_ACCESS_KEY
-)
-
-if "page" not in st.session_state:
-    st.session_state.page = "home"
-def upload_csv_file(file, csv_filename, content_type):
-    try:
-        # Generate a unique key for the file using UUID
-        uuid_str = str(uuid.uuid4())
-        s3_key = f'MoSPI_csv_files/{uuid_str}-{csv_filename}'
-
-        # Upload the CSV to S3
-        s3.upload_fileobj(
-            file,
-            AWS_BUCKET_NAME,
-            s3_key,
-            ExtraArgs={'ContentType': content_type}  # Set the MIME type of the uploaded file
-        )
-
-        upload_time = datetime.now()
-
-        # Metadata for MongoDB
-        metadata = {
-            'name': csv_filename,
-            'type': content_type,
-            's3_url': f's3://{AWS_BUCKET_NAME}/{s3_key}',
-            's3_key': s3_key,
-            'object_url': f'https://{AWS_BUCKET_NAME}.s3.amazonaws.com/{s3_key}',
-            'date_uploaded': upload_time.strftime('%Y-%m-%d'),
-            'time_uploaded': upload_time.strftime('%H:%M:%S')
-        }
-
-        return metadata
-
-    except Exception as e:
-        print(f"An error occurred during upload: {e}")
-        return None
-
-def process_image(url, filename):
-    try:
-
-        image_data = base64.b64encode(httpx.get(url).content).decode("utf-8")
-        if image_data:
-
-            result = process_image_using_llm(image_data, 1, 3)
-            has_table_data = result.get("has_table_data")
-            if has_table_data:
-                table_data = result.get("table_data")
-                page_number = result.get("page_number")
-                description = result.get("description")
-                column_summary = result.get("column_summary")
-                best_col1=result.get("best_col1")
-                best_col2=result.get("best_col2")
-
-                data={
-                    "table_data":table_data,
-                    "page_number":page_number,
-                    "description":description,
-                    "column_summary":column_summary,
-                    "best_col1":best_col1,
-                    "best_col2":best_col2
-                }
-
-
-                collection.update_one({"object_url": url}, {"$set": {"table_data": data}})
-                print("Successfully extracted data from image and inserted into MongoDB")
-
-                # Generate CSV from table data
-                csv_buffer = StringIO()
-                csv_writer = csv.DictWriter(csv_buffer, fieldnames=table_data[0].keys())
-                csv_writer.writeheader()
-                csv_writer.writerows(table_data)
-
-                # Convert CSV text to bytes for uploading
-                csv_bytes = BytesIO(csv_buffer.getvalue().encode("utf-8"))
-
-                # Upload CSV to S3
-                csv_filename = f"{filename}.csv"
-                s3_metadata = upload_csv_file(csv_bytes, csv_filename, content_type="text/csv")
-
-                if s3_metadata:
-                    # Update MongoDB with CSV S3 URL
-                    collection.update_one(
-                        {"object_url": url},
-                        {"$set": {
-                            "csv_object_url": s3_metadata.get("object_url"),
-                            "csv_s3_url": s3_metadata.get("s3_url")
-                        }}
-                    )
-                    print("CSV file uploaded to S3 and URL saved in MongoDB")
-
-                return True
-            else:
-                print(f"No table data was found in the image {url}")
-                return False
-
-        else:
-            print(f"No image data found in uploaded image")
-            return False
-
-    except Exception as e:
-        print(f"Error occurred in processing image: {e}")
-        return False
-
-
-def convert_excel_to_csv(file, filename):
-    # Determine the appropriate engine based on file extension
-    file_extension = filename.split('.')[-1].lower()
-    if file_extension == 'xlsx':
-        engine = 'openpyxl'
-    elif file_extension == 'xls':
-        engine = 'xlrd'
-    else:
-        raise ValueError("Unsupported file format for Excel. Please upload an .xls or .xlsx file.")
-
-    # Load the Excel file into a DataFrame
-    df = pd.read_excel(file, engine=engine)
-
-    # Convert the DataFrame to CSV format in memory
-    csv_buffer = BytesIO()
-    df.to_csv(csv_buffer, index=False)
-    csv_buffer.seek(0)  # Move to the start of the buffer
-
-    # Generate a new filename for CSV
-    csv_filename = filename.replace(".xlsx", ".csv").replace(".xls", ".csv")
-    return csv_buffer, csv_filename
-
-
-
-if st.session_state.page=="home":
-    st.title("Smart Data Extraction and Analysis")
-
-    uploaded_file = st.file_uploader(
-        "Upload a file",
-        type=["png", "jpg", "jpeg", "pdf", "xlsx", "xls", "csv"],
-        accept_multiple_files=False,
-        help="Please upload only one file of type image, PDF, Excel, or CSV."
-    )
-
-    if uploaded_file and st.button("Upload"):
-        with st.spinner("Processing your file"):
-            file_copy = BytesIO(uploaded_file.getvalue())
-            file_type = uploaded_file.type
-            metadata = upload_file(uploaded_file, file_type)
-            if metadata:
-
-                object_url = metadata.get("object_url")
-                filename = metadata.get("name")
-
-                if "image" in file_type:  # Process image files
-                    processed = process_image(object_url, filename)
-                    if processed:
-                        collection.update_one({"object_url": object_url}, {"$set": {"status": "processed"}})
-                        st.success("Image processed and CSV file uploaded to S3 successfully.")
-                    else:
-                        collection.update_one({"object_url":object_url},{"$set":{"status":"failed"}})
-                        st.error("Error occured in processing Image, please try again later")
-
-                elif "pdf" in file_type:
-                    processed=process_pdf(object_url,filename)
-                    if processed:
-                        collection.update_one({"object_url": object_url}, {"$set": {"status": "processed"}})
-                        st.success("Successfully processed pdf")
-                    else:
-                        collection.update_one({"object_url": object_url}, {"$set": {"status": "failed"}})
-                        st.error("Error occured in processing pdf")
-
-                elif file_type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet","application/vnd.ms-excel"]:
-                    csv_buffer, csv_filename = convert_excel_to_csv(file_copy, filename)
-                    s3_metadata = upload_csv_file(csv_buffer, csv_filename, content_type="text/csv")
-                    if s3_metadata:
-                        collection.update_one({"object_url": object_url}, {
-                            "$set": {"csv_object_url": s3_metadata["object_url"], "csv_s3_url": s3_metadata["s3_url"],
-                                     "filetype": "excel","status":"processed"}
-                        })
-                        st.success("Excel file uploaded to S3 successfully.")
-
-                    else:
-                        collection.update_one({"object_url": object_url}, {"$set": {"status": "failed"}})
-
-                elif "csv" in file_type:
-                    collection.update_one({"object_url": object_url}, {
-                        "$set": {"csv_object_url": object_url,"filetype": "csv","status":"processed"}})
-                    st.success("CSV file uploaded to S3 successfully.")
-
-    st.markdown("<hr>",unsafe_allow_html=True)
-    col1, col2, col3 = st.columns([1, 1, 1], gap="small")
-
-    with col1:
-        if st.button("View PDFs", key="View pdf button"):
-            st.session_state.page = "view_pdf"
-            st.rerun()
-
-    with col2:
-        if st.button("View Images", key="View image button"):
-            st.session_state.page = "view_image"
-            st.rerun()
-
-    with col3:
-        if st.button("View Excel", key="View excel button"):
-            st.session_state.page = "view_excel"
-            st.rerun()
-
-            #in case of csv we are already uploading it.
-
-if st.session_state.page=="view_pdf":
-    view_pdfs()
-
-
-elif st.session_state.page=="view_image":
-    view_images()
-
-elif st.session_state.page=="view_excel":
-    view_excel()
-
-if st.session_state.page=="view_image_analysis" and "image_url" in st.session_state:
-    image_url = st.session_state.image_url
-    view_table_analysis_page(image_url)
-
-if st.session_state.page=="pdf_analysis" and "pdf_url" in st.session_state:
-    pdf_url=st.session_state.pdf_url
-    view_pdf_table_analysis_page(pdf_url)
-
-
-if st.session_state.page=="view_excel_analysis" and "excel_url" in st.session_state:
-    excel_url=st.session_state.excel_url
-    display_csv_analysis(excel_url)
-
-
-
-
-
-
-
-
+import streamlit as st
+from upload_file_to_s3 import upload_file
+import base64
+import httpx
+from extract_table_from_image import process_image_using_llm
+from process_pdf import process_pdf
+from pymongo import MongoClient
+from datetime import datetime
+from table_analysis_for_image import view_table_analysis_page
+from table_analysis_for_pdf import view_pdf_table_analysis_page
+from table_analysis_for_excel import display_csv_analysis
+from view_excel import view_excel
+from copy import deepcopy
+
+import uuid
+import os
+import csv
+from view_pdf import view_pdfs
+from view_image import view_images
+from io import StringIO, BytesIO
+from dotenv import load_dotenv
+import boto3
+import pandas as pd
+st.set_page_config(layout='wide',page_title="MoSPI", page_icon="📄")
+load_dotenv()
+
+AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
+AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
+AWS_BUCKET_NAME = os.getenv("AWS_BUCKET_NAME")
+MONGO_URI = os.getenv("MONGO_URI")
+DB_NAME = os.getenv("DB_NAME")
+COLLECTION_NAME = os.getenv("COLLECTION_NAME")
+
+mongo_client = MongoClient(MONGO_URI)
+db = mongo_client[DB_NAME]
+collection = db[COLLECTION_NAME]
+
+s3 = boto3.client(
+    's3',
+    aws_access_key_id=AWS_ACCESS_KEY_ID,
+    aws_secret_access_key=AWS_SECRET_ACCESS_KEY
+)
+path_to_logo='logo.png'
+if "page" not in st.session_state:
+    st.session_state.page = "home"
+def upload_csv_file(file, csv_filename, content_type):
+    try:
+        # Generate a unique key for the file using UUID
+        uuid_str = str(uuid.uuid4())
+        s3_key = f'MoSPI_csv_files/{uuid_str}-{csv_filename}'
+
+        # Upload the CSV to S3
+        s3.upload_fileobj(
+            file,
+            AWS_BUCKET_NAME,
+            s3_key,
+            ExtraArgs={'ContentType': content_type}  # Set the MIME type of the uploaded file
+        )
+
+        upload_time = datetime.now()
+
+        # Metadata for MongoDB
+        metadata = {
+            'name': csv_filename,
+            'type': content_type,
+            's3_url': f's3://{AWS_BUCKET_NAME}/{s3_key}',
+            's3_key': s3_key,
+            'object_url': f'https://{AWS_BUCKET_NAME}.s3.amazonaws.com/{s3_key}',
+            'date_uploaded': upload_time.strftime('%Y-%m-%d'),
+            'time_uploaded': upload_time.strftime('%H:%M:%S')
+        }
+
+        return metadata
+
+    except Exception as e:
+        print(f"An error occurred during upload: {e}")
+        return None
+
+def process_image(url, filename):
+    try:
+
+        image_data = base64.b64encode(httpx.get(url).content).decode("utf-8")
+        if image_data:
+
+            result = process_image_using_llm(image_data, 1, 3)
+            has_table_data = result.get("has_table_data")
+            if has_table_data:
+                table_data = result.get("table_data")
+                page_number = result.get("page_number")
+                description = result.get("description")
+                column_summary = result.get("column_summary")
+                best_col1=result.get("best_col1")
+                best_col2=result.get("best_col2")
+
+                data={
+                    "table_data":table_data,
+                    "page_number":page_number,
+                    "description":description,
+                    "column_summary":column_summary,
+                    "best_col1":best_col1,
+                    "best_col2":best_col2
+                }
+
+
+                collection.update_one({"object_url": url}, {"$set": {"table_data": data}})
+                print("Successfully extracted data from image and inserted into MongoDB")
+
+                # Generate CSV from table data
+                csv_buffer = StringIO()
+                csv_writer = csv.DictWriter(csv_buffer, fieldnames=table_data[0].keys())
+                csv_writer.writeheader()
+                csv_writer.writerows(table_data)
+
+                # Convert CSV text to bytes for uploading
+                csv_bytes = BytesIO(csv_buffer.getvalue().encode("utf-8"))
+
+                # Upload CSV to S3
+                csv_filename = f"{filename}.csv"
+                s3_metadata = upload_csv_file(csv_bytes, csv_filename, content_type="text/csv")
+
+                if s3_metadata:
+                    # Update MongoDB with CSV S3 URL
+                    collection.update_one(
+                        {"object_url": url},
+                        {"$set": {
+                            "csv_object_url": s3_metadata.get("object_url"),
+                            "csv_s3_url": s3_metadata.get("s3_url")
+                        }}
+                    )
+                    print("CSV file uploaded to S3 and URL saved in MongoDB")
+
+                return True
+            else:
+                print(f"No table data was found in the image {url}")
+                return False
+
+        else:
+            print(f"No image data found in uploaded image")
+            return False
+
+    except Exception as e:
+        print(f"Error occurred in processing image: {e}")
+        return False
+
+
+def convert_excel_to_csv(file, filename):
+    # Determine the appropriate engine based on file extension
+    file_extension = filename.split('.')[-1].lower()
+    if file_extension == 'xlsx':
+        engine = 'openpyxl'
+    elif file_extension == 'xls':
+        engine = 'xlrd'
+    else:
+        raise ValueError("Unsupported file format for Excel. Please upload an .xls or .xlsx file.")
+
+    # Load the Excel file into a DataFrame
+    df = pd.read_excel(file, engine=engine)
+
+    # Convert the DataFrame to CSV format in memory
+    csv_buffer = BytesIO()
+    df.to_csv(csv_buffer, index=False)
+    csv_buffer.seek(0)  # Move to the start of the buffer
+
+    # Generate a new filename for CSV
+    csv_filename = filename.replace(".xlsx", ".csv").replace(".xls", ".csv")
+    return csv_buffer, csv_filename
+
+
+
+if st.session_state.page=="home":
+
+    col1,col2=st.columns([1,7])
+    with col1:
+        st.image(path_to_logo, width=100)
+    with col2:
+        st.title("Smart Data Extraction and Analysis tool")
+
+    uploaded_file = st.file_uploader(
+        "Upload a file",
+        type=["png", "jpg", "jpeg", "pdf", "xlsx", "xls", "csv"],
+        accept_multiple_files=False,
+        help="Please upload only one file of type image, PDF, Excel, or CSV."
+    )
+
+    if uploaded_file and st.button("Upload"):
+        with st.spinner("Processing your file"):
+            file_copy = BytesIO(uploaded_file.getvalue())
+            file_type = uploaded_file.type
+            metadata = upload_file(uploaded_file, file_type)
+            if metadata:
+
+                object_url = metadata.get("object_url")
+                filename = metadata.get("name")
+
+                if "image" in file_type:  # Process image files
+                    processed = process_image(object_url, filename)
+                    if processed:
+                        collection.update_one({"object_url": object_url}, {"$set": {"status": "processed"}})
+                        st.success("Image processed and CSV file uploaded to S3 successfully.")
+                    else:
+                        collection.update_one({"object_url":object_url},{"$set":{"status":"failed"}})
+                        st.error("Error occured in processing Image, please try again later")
+
+                elif "pdf" in file_type:
+                    processed=process_pdf(object_url,filename)
+                    if processed:
+                        collection.update_one({"object_url": object_url}, {"$set": {"status": "processed"}})
+                        st.success("Successfully processed pdf")
+                    else:
+                        collection.update_one({"object_url": object_url}, {"$set": {"status": "failed"}})
+                        st.error("Error occured in processing pdf")
+
+                elif file_type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet","application/vnd.ms-excel"]:
+                    csv_buffer, csv_filename = convert_excel_to_csv(file_copy, filename)
+                    s3_metadata = upload_csv_file(csv_buffer, csv_filename, content_type="text/csv")
+                    if s3_metadata:
+                        collection.update_one({"object_url": object_url}, {
+                            "$set": {"csv_object_url": s3_metadata["object_url"], "csv_s3_url": s3_metadata["s3_url"],
+                                     "filetype": "excel","status":"processed"}
+                        })
+                        st.success("Excel file uploaded to S3 successfully.")
+
+                    else:
+                        collection.update_one({"object_url": object_url}, {"$set": {"status": "failed"}})
+
+                elif "csv" in file_type:
+                    collection.update_one({"object_url": object_url}, {
+                        "$set": {"csv_object_url": object_url,"filetype": "csv","status":"processed"}})
+                    st.success("CSV file uploaded to S3 successfully.")
+
+    st.markdown("<hr>",unsafe_allow_html=True)
+    col1, col2, col3 = st.columns([1, 1, 1], gap="small")
+
+    with col1:
+        if st.button("View PDFs", key="View pdf button"):
+            st.session_state.page = "view_pdf"
+            st.rerun()
+
+    with col2:
+        if st.button("View Images", key="View image button"):
+            st.session_state.page = "view_image"
+            st.rerun()
+
+    with col3:
+        if st.button("View Excel", key="View excel button"):
+            st.session_state.page = "view_excel"
+            st.rerun()
+
+            #in case of csv we are already uploading it.
+
+if st.session_state.page=="view_pdf":
+    view_pdfs()
+
+
+elif st.session_state.page=="view_image":
+    view_images()
+
+elif st.session_state.page=="view_excel":
+    view_excel()
+
+if st.session_state.page=="view_image_analysis" and "image_url" in st.session_state:
+    image_url = st.session_state.image_url
+    view_table_analysis_page(image_url)
+
+if st.session_state.page=="pdf_analysis" and "pdf_url" in st.session_state:
+    pdf_url=st.session_state.pdf_url
+    view_pdf_table_analysis_page(pdf_url)
+
+
+if st.session_state.page=="view_excel_analysis" and "excel_url" in st.session_state:
+    excel_url=st.session_state.excel_url
+    display_csv_analysis(excel_url)
+
+
+
+
+
+
+
+