Update upload_image_page.py

upload_image_page.py

@@ -1,88 +1,104 @@
-import streamlit as st
-from pymongo import MongoClient
-import os
-from dotenv import load_dotenv
-from helper.upload_file_to_s3 import upload_file
-from helper.process_image import process_image_using_llm
-from helper.create_embeddings import create_embedding
-import time
-
-# Load environment variables
-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")
-COLLECTION_NAME2=os.getenv("COMPANY_COLLECTION_NAME")
-
-mongo_client = MongoClient(MONGO_URI)
-db = mongo_client[DB_NAME]
-collection = db[COLLECTION_NAME]
-collection2=db[COLLECTION_NAME2]
-
-def upload():
-    if st.button("Back"):
-        st.session_state.page = "upload_main"
-        st.rerun()
-
-    # File uploader (image files only)
-    uploaded_image = st.file_uploader("Choose an image file to upload", type=["png", "jpg", "jpeg"],
-                                      accept_multiple_files=False)
-
-    # Fetch tags and categories from MongoDB
-    tags_doc = collection2.find_one({"type": "tags"})
-    categories_doc = collection2.find_one({"type": "categories"})
-
-    tags_options = tags_doc["tags"] if tags_doc and "tags" in tags_doc else []
-    categories_options = categories_doc["categories"] if categories_doc and "categories" in categories_doc else []
-
-    # Multi-select dropdowns for tags and categories
-    selected_tags = st.multiselect("Select Tags", options=tags_options)
-    selected_categories = st.multiselect("Select Categories", options=categories_options)
-
-    if uploaded_image and selected_tags and selected_categories:
-        flag=False
-        if st.button("Submit"):
-
-            with st.spinner(text="Uploading and Processing Image"):
-                # Upload file to S3
-                metadata = upload_file(uploaded_image,"Image")
-                if metadata:
-                    object_url = metadata.get("object_url")
-                    filename = metadata.get("name")
-
-                    # Process image with LLM for description
-                    llm_processed = process_image_using_llm(object_url)
-                    if llm_processed:
-                        # Create embedding with tags and categories in metadata
-                        embedding_created = create_embedding(
-                            object_url,
-                            selected_tags,
-                            selected_categories
-                        )
-                        if embedding_created:
-                            # Save tags and categories to MongoDB document for the uploaded image
-                            collection.update_one(
-                                {"object_url": object_url},
-                                {"$set": {
-                                    "tags": selected_tags,
-                                    "categories": selected_categories
-                                }}
-                            )
-                            st.success("Image has been successfully uploaded and processed.")
-                            flag=True
-                        else:
-                            st.error("Could not create embedding. Please try again.")
-                    else:
-                        st.error("Could not process the image description. Please try again.")
-                else:
-                    st.error("Could not upload the image. Please try again.")
-
-
-            if flag:
-                st.write("Redirecting to View Page to view all uploaded images")
-                time.sleep(2)
-                st.session_state.page = "view_image"
-                st.rerun()
+import streamlit as st
+from pymongo import MongoClient
+import os
+from dotenv import load_dotenv
+from helper.upload_file_to_s3 import upload_file
+from helper.process_image import process_image_using_llm
+from helper.create_embeddings import create_embedding
+import time
+
+# Load environment variables
+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")
+COLLECTION_NAME2=os.getenv("COMPANY_COLLECTION_NAME")
+
+mongo_client = MongoClient(MONGO_URI)
+db = mongo_client[DB_NAME]
+collection = db[COLLECTION_NAME]
+collection2=db[COLLECTION_NAME2]
+
+def upload():
+    if st.button("Back"):
+        st.session_state.page = "upload_main"
+        st.rerun()
+
+    # File uploader (image files only)
+    uploaded_image = st.file_uploader("Choose an image file to upload", type=["png", "jpg", "jpeg"],
+                                      accept_multiple_files=False)
+
+    # Fetch tags and categories from MongoDB
+    tags_doc = collection2.find_one({"type": "tags"})
+    categories_doc = collection2.find_one({"type": "categories"})
+
+    tags_options = tags_doc["tags"] if tags_doc and "tags" in tags_doc else []
+    categories_options = categories_doc["categories"] if categories_doc and "categories" in categories_doc else []
+
+    # Multi-select dropdowns for tags and categories
+    selected_tags = st.multiselect("Select Tags", options=tags_options)
+    selected_categories = st.multiselect("Select Categories", options=categories_options)
+
+    if uploaded_image and selected_tags and selected_categories:
+        flag=False
+        if st.button("Submit"):
+
+            with st.spinner(text="Uploading and Processing Image"):
+                # Upload file to S3
+                metadata = upload_file(uploaded_image,"Image")
+                if metadata:
+                    object_url = metadata.get("object_url")
+                    filename = metadata.get("name")
+
+                    # Process image with LLM for description
+                    llm_processed = process_image_using_llm(object_url)
+                    if llm_processed:
+                        # Create embedding with tags and categories in metadata
+                        embedding_created = create_embedding(
+                            object_url,
+                            selected_tags,
+                            selected_categories
+                        )
+                        if embedding_created:
+                            # Save tags and categories to MongoDB document for the uploaded image
+                            collection.update_one(
+                                {"object_url": object_url},
+                                {"$set": {
+                                    "tags": selected_tags,
+                                    "categories": selected_categories,
+                                    "status":"processed"
+                                }}
+                            )
+                            st.success("Image has been successfully uploaded and processed.")
+                            flag=True
+                        else:
+                            st.error("Could not create embedding. Please try again.")
+                            collection.update_one(
+                                {"object_url": object_url},
+                                {"$set": {
+
+                                    "status": "failed"
+                                }}
+                            )
+                    else:
+                        st.error("Could not process the image description. Please try again.")
+                        st.error("Could not create embedding. Please try again.")
+                        collection.update_one(
+                            {"object_url": object_url},
+                            {"$set": {
+
+                                "status": "failed"
+                            }}
+                        )
+                else:
+                    st.error("Could not upload the image. Please try again.")
+
+
+            if flag:
+                st.write("Redirecting to View Page to view all uploaded images")
+                time.sleep(2)
+                st.session_state.page = "view_image"
+                st.rerun()