app.py

import os

import streamlit as st
from dotenv import load_dotenv
from PIL import Image

from blip_model import extract_image_details
from qwen_model import generate_response

# Load environment variables from .env (if you have any relevant ones)
load_dotenv()

# Initialize session state for image details
if "image_details" not in st.session_state:
    st.session_state.image_details = ""

st.title("Image to Text Response with RAG (Qwen)")

# File uploader for image upload
uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Check if a new file is uploaded
    if "uploaded_file" not in st.session_state or uploaded_file != st.session_state.uploaded_file:
        st.session_state.uploaded_file = uploaded_file

        # Clear previous messages
        st.session_state.messages = []

        # Show loading status while processing the image
        with st.status("Processing image...", state="running"):
            image = Image.open(uploaded_file)
            # Extract text details from the image (using your BLIP model)
            st.session_state.image_details = extract_image_details(image)

        st.success("Image processed successfully.")
else:
    # If no file is uploaded, stop the execution
    st.stop()

# Chat interface
if st.session_state.image_details:
    st.write("### Ask a question about the image")
    
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    # Chat input
    prompt = st.chat_input("Ask something about the image")
    if prompt:
        st.session_state.messages.append({"role": "user", "content": prompt})
    
        # Generate response
        response = generate_response([st.session_state.image_details], prompt)
        st.session_state.messages.append({"role": "assistant", "content": response})
    
        # Display user message
        with st.chat_message("user"):
            st.write(prompt)
    
        # Display assistant message
        with st.chat_message("assistant"):
            st.write(response)