app.py

import os
import zipfile
import time
import streamlit as st
from PIL import Image as PILImage
from PIL import Image as pilImage
import base64
import io
import chromadb
from initate import process_pdf
from utils.llm_ag import intiate_convo
from utils.doi import process_image_and_get_description

path = "mm_vdb2"
client = chromadb.PersistentClient(path=path)
import streamlit as st
from PIL import Image as PILImage

def display_images(image_collection, query_text, max_distance=None, debug=False):
    """
    Display images in a Streamlit app based on a query.

    Args:
        image_collection: The image collection object for querying.
        query_text (str): The text query for images.
        max_distance (float, optional): Maximum allowable distance for filtering.
        debug (bool, optional): Whether to print debug information.
    """
    results = image_collection.query(
        query_texts=[query_text],
        n_results=10,
        include=['uris', 'distances']
    )

    uris = results['uris'][0]
    distances = results['distances'][0]

    # Combine uris and distances, then sort by URI in ascending order
    sorted_results = sorted(zip(uris, distances), key=lambda x: x[0])

    # Filter and display images
    for uri, distance in sorted_results:
        if max_distance is None or distance <= max_distance:
            if debug:
                st.write(f"URI: {uri} - Distance: {distance}")
            try:
                img = PILImage.open(uri)
                st.image(img, width=300)
            except Exception as e:
                st.error(f"Error loading image {uri}: {e}")
        else:
            if debug:
                st.write(f"URI: {uri} - Distance: {distance} (Filtered out)")

            

def display_videos_streamlit(video_collection, query_text, max_distance=None, max_results=5, debug=False):
    """
    Display videos in a Streamlit app based on a query.

    Args:
        video_collection: The video collection object for querying.
        query_text (str): The text query for videos.
        max_distance (float, optional): Maximum allowable distance for filtering.
        max_results (int, optional): Maximum number of results to display.
        debug (bool, optional): Whether to print debug information.
    """
    # Deduplication set
    displayed_videos = set()

    # Query the video collection with the specified text
    results = video_collection.query(
        query_texts=[query_text],
        n_results=max_results,  # Adjust the number of results if needed
        include=['uris', 'distances', 'metadatas']
    )

    # Extract URIs, distances, and metadatas from the result
    uris = results['uris'][0]
    distances = results['distances'][0]
    metadatas = results['metadatas'][0]

    # Display the videos that meet the distance criteria
    for uri, distance, metadata in zip(uris, distances, metadatas):
        video_uri = metadata['video_uri']

        # Check if a max_distance filter is applied and the distance is within the allowed range
        if (max_distance is None or distance <= max_distance) and video_uri not in displayed_videos:
            if debug:
                st.write(f"URI: {uri} - Video URI: {video_uri} - Distance: {distance}")
            st.video(video_uri)  # Display video in Streamlit
            displayed_videos.add(video_uri)  # Add to the set to prevent duplication
        else:
            if debug:
                st.write(f"URI: {uri} - Video URI: {video_uri} - Distance: {distance} (Filtered out)")
                
                
def image_uris(image_collection,query_text, max_distance=None, max_results=5):
    results = image_collection.query(
        query_texts=[query_text],
        n_results=max_results,
        include=['uris', 'distances']
    )

    filtered_uris = []
    for uri, distance in zip(results['uris'][0], results['distances'][0]):
        if max_distance is None or distance <= max_distance:
            filtered_uris.append(uri)

    return filtered_uris

def text_uris(text_collection,query_text, max_distance=None, max_results=5):
    results = text_collection.query(
        query_texts=[query_text],
        n_results=max_results,
        include=['documents', 'distances']
    )

    filtered_texts = []
    for doc, distance in zip(results['documents'][0], results['distances'][0]):
        if max_distance is None or distance <= max_distance:
            filtered_texts.append(doc)

    return filtered_texts

def frame_uris(video_collection,query_text, max_distance=None, max_results=5):
    results = video_collection.query(
        query_texts=[query_text],
        n_results=max_results,
        include=['uris', 'distances']
    )

    filtered_uris = []
    seen_folders = set()

    for uri, distance in zip(results['uris'][0], results['distances'][0]):
        if max_distance is None or distance <= max_distance:
            folder = os.path.dirname(uri)
            if folder not in seen_folders:
                filtered_uris.append(uri)
                seen_folders.add(folder)

        if len(filtered_uris) == max_results:
            break

    return filtered_uris

def image_uris2(image_collection2,query_text, max_distance=None, max_results=5):
    results = image_collection2.query(
        query_texts=[query_text],
        n_results=max_results,
        include=['uris', 'distances']
    )

    filtered_uris = []
    for uri, distance in zip(results['uris'][0], results['distances'][0]):
        if max_distance is None or distance <= max_distance:
            filtered_uris.append(uri)

    return filtered_uris


def format_prompt_inputs(image_collection, text_collection, video_collection, user_query):
    frame_candidates = frame_uris(video_collection, user_query, max_distance=1.55)
    image_candidates = image_uris(image_collection, user_query, max_distance=1.5)
    texts = text_uris(text_collection, user_query, max_distance=1.3)

    inputs = {"query": user_query, "texts": texts}
    frame = frame_candidates[0] if frame_candidates else ""
    inputs["frame"] = frame

    if image_candidates:
        image = image_candidates[0]
        with PILImage.open(image) as img:
            img = img.resize((img.width // 2, img.height // 2))
            # img = img.convert("L")
            with io.BytesIO() as output:
                img.save(output, format="JPEG", quality=60)
                compressed_image_data = output.getvalue()

        inputs["image_data_1"] = base64.b64encode(compressed_image_data).decode('utf-8')
    else:
        inputs["image_data_1"] = ""

    return inputs
    
def unzip_file(zip_path, extract_to):
    """
    Unzips a zip file to the specified directory.
    
    Args:
        zip_path (str): Path to the zip file.
        extract_to (str): Directory where the contents should be extracted.
    """
    try:
        # Ensure the destination directory exists
        os.makedirs(extract_to, exist_ok=True)

        # Open the zip file
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            # Extract all the contents
            zip_ref.extractall(extract_to)
        return True
    except Exception as e:
        print(f"An error occurred: {e}")
        return False

def process_pdf(pdf_path):
    # Placeholder function to simulate PDF processing
    # Replace this with actual PDF processing logic, such as extracting text, images, etc.
    time.sleep(2)  # Simulating processing delay
    return "image_collection", "text_collection", "video_collection"  # Replace with actual collections

def page_1():
    st.title("Page 1: Upload and Process Videos and PDFs")

    # File uploader for multiple zip files containing videos
    uploaded_video_zips = st.file_uploader("Upload ZIP files containing videos", type=["zip"], accept_multiple_files=True)

    # File uploader for PDF files
    uploaded_pdf_files = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)

    # Button to trigger processing
    if (uploaded_video_zips or uploaded_pdf_files) and st.button("Process Files"):
        # Temporary folder to store extracted files
        temp_folder = "/tmp/extracted_files"
        os.makedirs(temp_folder, exist_ok=True)

        # Progress bar
        progress_bar = st.progress(0)
        status_text = st.empty()

        try:
            total_files = len(uploaded_video_zips) + len(uploaded_pdf_files)
            files_processed = 0
            progress_step = 100 / total_files if total_files > 0 else 0

            # Process video zip files
            for uploaded_file in uploaded_video_zips:
                zip_path = f"/tmp/{uploaded_file.name}"
                with open(zip_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())

                # Extract the content from the zip file
                folder_name = os.path.splitext(uploaded_file.name)[0]
                extract_to = os.path.join(temp_folder, folder_name)
                if unzip_file(zip_path, extract_to):
                    files_processed += 1
                    progress_bar.progress(files_processed * progress_step)
                    status_text.text(f"Extracting: {uploaded_file.name} ({files_processed}/{total_files})")

            # Process PDF files
            for uploaded_pdf in uploaded_pdf_files:
                pdf_path = f"/tmp/{uploaded_pdf.name}"
                with open(pdf_path, "wb") as f:
                    f.write(uploaded_pdf.getbuffer())

                # Simulate PDF processing (replace with actual PDF processing logic)
                files_processed += 1
                progress_bar.progress(files_processed * progress_step)
                status_text.text(f"Processing PDF: {uploaded_pdf.name} ({files_processed}/{total_files})")
                
                # Call your actual PDF processing function here, e.g.
                image_collection, text_collection, video_collection = process_pdf(pdf_path,temp_folder)

                # Save collections to session state
                st.session_state.image_collection = image_collection
                st.session_state.text_collection = text_collection
                st.session_state.video_collection = video_collection

            # Update status after extraction and processing
            status_text.text("Extraction and processing completed successfully!")

            st.success("Videos and PDFs processed successfully! Collections saved to session state.")

        except Exception as e:
            progress_bar.progress(0)
            status_text.text("")
            st.error(f"Error processing files: {e}")

def page_2():
    st.title("Page 2: Query and Use Processed Collections")

    if "image_collection" in st.session_state and "text_collection" in st.session_state and "video_collection" in st.session_state:
        image_collection = st.session_state.image_collection
        text_collection = st.session_state.text_collection
        video_collection = st.session_state.video_collection
        st.success("Collections loaded successfully.")

        query = st.text_input("Enter your query", value="Example Query")
        if query:
            inputs = format_prompt_inputs(image_collection, text_collection, video_collection, query)
            texts = inputs["texts"]
            image_data_1 = inputs["image_data_1"]

            if image_data_1:
                image_data_1 = process_image_and_get_description(image_data_1)

            response = intiate_convo(query, image_data_1, texts)
            st.write("Response:", response)

            st.markdown("### Images")
            display_images(image_collection, query, max_distance=1.55, debug=True)

            st.markdown("### Videos")
            display_videos_streamlit(video_collection, query_text=query, max_distance=None, max_results=5, debug=False)
            frame = inputs["frame"]
            
            if frame:
                video_path = f"video/StockVideos-CC0/{os.path.basename(frame).split('/')[0]}.mp4"
                if os.path.exists(video_path):
                    st.video(video_path)
                else:
                    st.write("No related videos found.")
    else:
        st.error("Collections not found in session state. Please process the PDF on Page 1.")

# --- Navigation ---

PAGES = {
    "Upload and Process PDF": page_1,
    "Query and Use Processed Collections": page_2
}

# Select page
selected_page = st.sidebar.selectbox("Choose a page", options=list(PAGES.keys()))

# Render selected page
PAGES[selected_page]()