app.py

import streamlit as st
from PIL import Image
from main import ImageProcessor, PDFProcessor  # Replace 'your_module' with the actual module where your ImageProcessor and PDFProcessor classes are defined
import google.generativeai as genai
import os

genai.configure(api_key=os.getenv("API_KEY"))
llm = genai.GenerativeModel("gemini-pro")


# Streamlit title
st.title("Q&A with Images & PDF")

# File upload for image or PDF
file_type = st.radio("Select File Type:", ["Image", "PDF"])
uploaded_file = st.file_uploader("Upload file:", type=["jpg", "jpeg", "png", "pdf"])

# Query text box
query = st.text_input("Enter your question:")

# Display box for generated answer
answer_display = st.empty()

if uploaded_file is not None:
    if file_type == "Image":
        # Image processing
        image_processor = ImageProcessor(uploaded_file)
        captions = image_processor.get_caption(uploaded_file)
        detections = image_processor.detect_objects(uploaded_file)
        prompt = image_processor.make_prompt(query, captions, detections)
        answer = image_processor.generate_answer(prompt)
        answer_display.text(answer)

    elif file_type == "PDF":
        # PDF processing
        pdf_processor = PDFProcessor(uploaded_file)
        pdf_vector_stores = pdf_processor.create_embedding_df(uploaded_file)
        relevant_passage = pdf_processor.find_best_passage(query, pdf_vector_stores)
        prompt = pdf_processor.make_prompt(query, relevant_passage)
        answer = pdf_processor.generate_answer(prompt)
        answer_display.text(answer)