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
from qdrant_client import QdrantClient, models
import getpass

curr_user = getpass.getuser()
# from langchain.vectorstores import Chroma
# To connect to the same event-loop,
# allows async events to run on notebook

# import nest_asyncio

# nest_asyncio.apply()

from dotenv import load_dotenv
load_dotenv()


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 calc_diff(img1 , img2) :
    i1 = Image.open(img1)
    i2 = Image.open(img2)
    h1 = imagehash.phash(i1)
    h2 = imagehash.phash(i2)
    return h1 - h2

def remove_duplicate_images(data_path) :
    image_files = os.listdir(data_path)
    only_images = []
    for one_image in image_files : 
        if one_image.endswith('jpeg') or one_image.endswith('png') or one_image.endswith('jpg') :
            only_images.append(one_image)
    only_images1 = sorted(only_images) 
    for one_image in only_images1 :
        for another_image in only_images1 :
            try :
                if one_image == another_image :
                    continue
                else :
                    diff = calc_diff(os.path.join(data_path ,one_image) , os.path.join(data_path ,another_image))
                    if diff ==0  :
                        os.remove(os.path.join(data_path , another_image))
            except Exception as e:
                print(e)
                pass
# from langchain_chroma import Chroma
# import chromadb
def initialize_qdrant(temp_dir , file_name , user):
    client = qdrant_client.QdrantClient(path=f"qdrant_mm_db_pipeline_{user}_{file_name}")
    # client = qdrant_client.QdrantClient(url = "http://localhost:2452")    
    # client = qdrant_client.QdrantClient(url="4b0af7be-d5b3-47ac-b215-128ebd6aa495.europe-west3-0.gcp.cloud.qdrant.io:6333", api_key="CO1sNGLmC6R_Q45qSIUxBSX8sxwHud4MCm4as_GTI-vzQqdUs-bXqw",)
    # client = qdrant_client.AsyncQdrantClient(location = ":memory:")
        
    if "vectordatabase" not in st.session_state or not st.session_state.vectordatabase:
            
        # text_store = client.create_collection(f"text_collection_pipeline_{user}_{file_name}"  ) 
        # image_store = client.create_collection(f"image_collection_pipeline_{user}_{file_name}"  ) 


        text_store = QdrantVectorStore( client = client , collection_name=f"text_collection_pipeline_{user}_{file_name}" )
        image_store = QdrantVectorStore(client = client , collection_name=f"image_collection_pipeline_{user}_{file_name}")
        storage_context = StorageContext.from_defaults(vector_store=text_store, image_store=image_store)
        documents = SimpleDirectoryReader(os.path.join(temp_dir, f"my_own_data_{user}_{file_name}")).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
#tmpnimvp35m , tmpnimvp35m
def process_pdf(pdf_file):
    temp_dir = tempfile.TemporaryDirectory()
    unique_folder_name = temp_dir.name.split('/')[-1]
    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, f"my_own_data_{unique_folder_name}_{os.path.splitext(pdf_file.name)[0]}")
    os.makedirs(data_path , exist_ok=True) 
    img_save_path = os.path.join(temp_dir.name, f"extracted_images_{unique_folder_name}_{os.path.splitext(pdf_file.name)[0]}")
    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)
    remove_duplicate_images(data_path)
    retriever_engine = initialize_qdrant(temp_dir.name , os.path.splitext(pdf_file.name)[0] , unique_folder_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!")

    if st.session_state.retriever_engine :
        query = st.text_input("Enter your question:")
        
        
        if st.button("Ask Question"):
            print("running")
            try:

                with st.spinner("Retrieving information..."):
                    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()