import streamlit as st
# Configure page
st.set_page_config(
page_title="E-commerce Visual Assistant",
page_icon="🛍️",
layout="wide"
)
from streamlit_chat import message
import torch
from PIL import Image
import requests
from io import BytesIO
from model import initialize_models, load_data, chatbot, cleanup_resources
# Helper functions
def load_image_from_url(url):
try:
response = requests.get(url)
img = Image.open(BytesIO(response.content))
return img
except Exception as e:
st.error(f"Error loading image from URL: {str(e)}")
return None
def initialize_assistant():
if not st.session_state.models_loaded:
with st.spinner("Loading models and data..."):
initialize_models()
load_data()
st.session_state.models_loaded = True
st.success("Assistant is ready!")
def display_chat_history():
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
if "image" in message:
st.image(message["image"], caption="Uploaded Image", width=200)
if "display_images" in message:
# Since we only have one image, we don't need multiple columns
img_data = message["display_images"][0] # Get the first (and only) image
st.image(
img_data['image'],
caption=f"{img_data['product_name']}\nPrice: ${img_data['price']:.2f}",
width=350 # Adjusted width for single image display
)
def handle_user_input(prompt, uploaded_image):
# Add user message
st.session_state.messages.append({"role": "user", "content": prompt})
# Generate response
with st.spinner("Processing your request..."):
try:
response = chatbot(prompt, image_input=uploaded_image)
if isinstance(response, dict):
assistant_message = {
"role": "assistant",
"content": response['text']
}
if 'images' in response and response['images']:
assistant_message["display_images"] = response['images']
st.session_state.messages.append(assistant_message)
else:
st.session_state.messages.append({
"role": "assistant",
"content": response
})
except Exception as e:
st.error(f"Error: {str(e)}")
st.session_state.messages.append({
"role": "assistant",
"content": f"I encountered an error: {str(e)}"
})
st.rerun()
# Custom CSS for enhanced styling
st.markdown("""
<style>
/* Main container styling */
.main {
background: linear-gradient(135deg, #f5f7fa 0%, #e8edf2 100%);
padding: 20px;
border-radius: 15px;
}
/* Header styling */
.stTitle {
color: #1e3d59;
font-size: 2.5rem !important;
text-align: center;
padding: 20px;
text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
}
/* Sidebar styling */
.css-1d391kg {
background: linear-gradient(180deg, #1e3d59 0%, #2b5876 100%);
}
/* Chat container styling */
.stChatMessage {
background-color: white;
border-radius: 15px;
box-shadow: 0 4px 6px rgba(0,0,0,0.1);
margin: 10px 0;
padding: 15px;
}
/* Input box styling */
.stTextInput > div > div > input {
border-radius: 20px;
border: 2px solid #1e3d59;
padding: 10px 20px;
}
/* Radio button styling */
.stRadio > label {
background-color: white;
padding: 10px 20px;
border-radius: 10px;
margin: 5px;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
/* Button styling */
.stButton > button {
background: linear-gradient(90deg, #1e3d59 0%, #2b5876 100%);
color: white;
border-radius: 20px;
padding: 10px 25px;
border: none;
box-shadow: 0 4px 6px rgba(0,0,0,0.1);
transition: all 0.3s ease;
}
.stButton > button:hover {
transform: translateY(-2px);
box-shadow: 0 6px 8px rgba(0,0,0,0.2);
}
/* Footer styling */
footer {
background-color: white;
border-radius: 10px;
padding: 20px;
margin-top: 30px;
text-align: center;
box-shadow: 0 4px 6px rgba(0,0,0,0.1);
}
</style>
""", unsafe_allow_html=True)
# Initialize session state
if 'messages' not in st.session_state:
st.session_state.messages = []
if 'models_loaded' not in st.session_state:
st.session_state.models_loaded = False
# Main title with enhanced styling
st.markdown("<h1 class='stTitle'>🛍️ Amazon E-commerce Visual Assistant</h1>", unsafe_allow_html=True)
# Sidebar configuration with enhanced styling
with st.sidebar:
st.title("Assistant Features")
st.markdown("### 🤖 How It Works")
st.markdown("""
This AI-powered shopping assistant combines:
**🧠 Advanced Technologies**
- FashionCLIP Visual AI
- Mistral-7B Language Model
- Multimodal Understanding
**💫 Capabilities**
- Product Search & Recognition
- Visual Analysis
- Detailed Comparisons
- Price Analysis
""")
st.markdown("---")
st.markdown("### 👥 Development Team")
team_members = {
"Yu-Chih (Wisdom) Chen",
"Feier Xu",
"Yanchen Dong",
"Kitae Kim"
}
for name in team_members:
st.markdown(f"**{name}**")
st.markdown("---")
if st.button("🔄 Reset Chat"):
st.session_state.messages = []
st.rerun()
# Main chat interface
def main():
# Initialize assistant
initialize_assistant()
# Chat container
chat_container = st.container()
# User input section at the bottom
input_container = st.container()
with input_container:
# Chat input
prompt = st.chat_input("What would you like to know?")
# Input options below chat input
col1, col2, col3 = st.columns([1,1,1])
with col1:
input_option = st.radio(
"Input Method:",
("Text Only", "Upload Image", "Image URL"),
key="input_method"
)
# Handle different input methods
uploaded_image = None
if input_option == "Upload Image":
with col2:
uploaded_file = st.file_uploader("Choose image", type=["jpg", "jpeg", "png"])
if uploaded_file:
uploaded_image = Image.open(uploaded_file)
st.image(uploaded_image, caption="Uploaded Image", width=200)
elif input_option == "Image URL":
with col2:
image_url = st.text_input("Enter image URL")
if image_url:
uploaded_image = load_image_from_url(image_url)
if uploaded_image:
st.image(uploaded_image, caption="Image from URL", width=200)
# Display chat history
with chat_container:
display_chat_history()
# Handle user input and generate response
if prompt:
handle_user_input(prompt, uploaded_image)
# Footer
st.markdown("""
<footer>
<h3>💡 Tips for Best Results</h3>
<p>Be specific in your questions for more accurate responses!</p>
<p>Try asking about product features, comparisons, or prices.</p>
</footer>
""", unsafe_allow_html=True)
if __name__ == "__main__":
try:
main()
finally:
cleanup_resources()