Spaces:
Sleeping
Sleeping
| 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() | |