Spaces:
Sleeping
Sleeping
File size: 4,501 Bytes
70124e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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()
|