File size: 1,563 Bytes
030c51b |
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 |
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)
|