app.py

import streamlit as st
import os
from PyPDF2 import PdfReader
import pymupdf
import numpy as np
import cv2
import shutil
import imageio
from PIL import Image
import imagehash
import matplotlib.pyplot as plt
from llama_index.core.indices import MultiModalVectorStoreIndex
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core import SimpleDirectoryReader, StorageContext
import qdrant_client
from llama_index.core import PromptTemplate
from llama_index.core.query_engine import SimpleMultiModalQueryEngine
from llama_index.llms.openai import OpenAI
from llama_index.core import load_index_from_storage, get_response_synthesizer
import tempfile


def extract_text_from_pdf(pdf_path):
    reader = PdfReader(pdf_path)
    full_text = ''
    for page in reader.pages:
        text = page.extract_text()
        full_text += text
    return full_text

def extract_images_from_pdf(pdf_path, img_save_path):
    doc = pymupdf.open(pdf_path)
    for page in doc:
        img_number = 0
        for block in page.get_text("dict")["blocks"]:
            if block['type'] == 1:
                name = os.path.join(img_save_path, f"img{page.number}-{img_number}.{block['ext']}")
                out = open(name, "wb")
                out.write(block["image"])
                out.close()
                img_number += 1

def is_empty(img_path):
    image = cv2.imread(img_path, 0)
    std_dev = np.std(image)
    return std_dev < 1

def move_images(source_folder, dest_folder):
    image_files = [f for f in os.listdir(source_folder) 
                   if f.lower().endswith(('.jpg', '.jpeg', '.png', '.gif'))]
    os.makedirs(dest_folder, exist_ok=True)
    moved_count = 0
    for file in image_files:
        src_path = os.path.join(source_folder, file)
        if not is_empty(src_path):
            shutil.move(src_path, os.path.join(dest_folder, file))
            moved_count += 1
    return moved_count

def remove_low_size_images(data_path):
    images_list = os.listdir(data_path)
    low_size_photo_list = []
    for one_image in images_list:
        image_path = os.path.join(data_path, one_image)
        try:
            pic = imageio.imread(image_path)
            size = pic.size
            if size < 100:
                low_size_photo_list.append(one_image)
        except:
            pass
    for one_image in low_size_photo_list[1:]:
        os.remove(os.path.join(data_path, one_image))

def initialize_qdrant(temp_dir):
    # try :
    client = qdrant_client.QdrantClient(path="qdrant_mm_db_pipeline")
    # except :
        # pass
    if "vectordatabase" not in st.session_state or not st.session_state.vectordatabase:
        text_store = QdrantVectorStore(client=client, collection_name="text_collection_pipeline")
        image_store = QdrantVectorStore(client=client, collection_name="image_collection_pipeline")
        storage_context = StorageContext.from_defaults(vector_store=text_store, image_store=image_store)
        documents = SimpleDirectoryReader(os.path.join(temp_dir, "my_own_data")).load_data()
        index = MultiModalVectorStoreIndex.from_documents(documents, storage_context=storage_context)
        st.session_state.vectordatabase = index
    else :
        index = st.session_state.vectordatabase
    retriever_engine = index.as_retriever(similarity_top_k=1, image_similarity_top_k=1)
    return retriever_engine

def plot_images(image_paths):
    images_shown = 0
    plt.figure(figsize=(16, 9))
    for img_path in image_paths:
        if os.path.isfile(img_path):
            image = Image.open(img_path)
            plt.subplot(2, 3, images_shown + 1)
            plt.imshow(image)
            plt.xticks([])
            plt.yticks([])
            images_shown += 1
            if images_shown >= 6:
                break

def retrieve_and_query(query, retriever_engine):
    retrieval_results = retriever_engine.retrieve(query)
    
    qa_tmpl_str = (
        "Context information is below.\n"
        "---------------------\n"
        "{context_str}\n"
        "---------------------\n"
        "Given the context information , "
        "answer the query in detail.\n"
        "Query: {query_str}\n"
        "Answer: "
    )
    qa_tmpl = PromptTemplate(qa_tmpl_str)

    llm = OpenAI(model="gpt-4o", temperature=0)
    response_synthesizer = get_response_synthesizer(response_mode="refine", text_qa_template=qa_tmpl, llm=llm)

    response = response_synthesizer.synthesize(query, nodes=retrieval_results)
    
    retrieved_image_path_list = []
    for node in retrieval_results:
        if (node.metadata['file_type'] == 'image/jpeg') or (node.metadata['file_type'] == 'image/png'):
            if node.score > 0.25:
                retrieved_image_path_list.append(node.metadata['file_path'])
    
    return response, retrieved_image_path_list

def process_pdf(pdf_file):
    # import pdb; pdb.set_trace()
    temp_dir = tempfile.TemporaryDirectory()
    temp_pdf_path = os.path.join(temp_dir.name, pdf_file.name)
    with open(temp_pdf_path, "wb") as f:
        f.write(pdf_file.getvalue())

    data_path = os.path.join(temp_dir.name, "my_own_data")
    os.makedirs(data_path , exist_ok=True) 
    img_save_path = os.path.join(temp_dir.name, "extracted_images")
    os.makedirs(img_save_path , exist_ok=True) 

    extracted_text = extract_text_from_pdf(temp_pdf_path)
    with open(os.path.join(data_path, "content.txt"), "w") as file:
        file.write(extracted_text)

    extract_images_from_pdf(temp_pdf_path, img_save_path)
    moved_count = move_images(img_save_path, data_path)
    remove_low_size_images(data_path)

    retriever_engine = initialize_qdrant(temp_dir.name)

    return temp_dir, retriever_engine

def main():
    st.title("PDF Vector Database Query Tool")
    st.markdown("This tool creates a vector database from a PDF and allows you to query it.")
    
    if "retriever_engine" not in st.session_state:
        st.session_state.retriever_engine = None
    if "vectordatabase" not in st.session_state:
        st.session_state.vectordatabase = None

    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
    if uploaded_file is None:
        st.info("Please upload a PDF file.")
    else:
        st.info(f"Uploaded PDF: {uploaded_file.name}")
        if st.button("Process PDF"):
            with st.spinner("Processing PDF..."):
                temp_dir, st.session_state.retriever_engine = process_pdf(uploaded_file)
        
                st.success("PDF processed successfully!")

    
    query = st.text_input("Enter your question:")
    
    
    if st.button("Ask Question"):
        print("running")
        try:
            import pdb; pdb.set_trace() 

            with st.spinner("Retrieving information..."):
                import pdb; pdb.set_trace()
                response, retrieved_image_path_list = retrieve_and_query(query, st.session_state.retriever_engine)
            
            st.write("Retrieved Context:")
            for node in response.source_nodes:
                st.code(node.node.get_text())
            
            st.write("\nRetrieved Images:")
            plot_images(retrieved_image_path_list)
            st.pyplot()

            st.write("\nFinal Answer:")
            st.code(response.response)
        
        except Exception as e:
            st.error(f"An error occurred: {e}")

if __name__ == "__main__":
    main()