Update user.py

user.py

@@ -17,55 +17,6 @@ def get_answer(query):
     #return response["result"]
     return response
 
-def home():
-    st.header("Welcome")
-    #st.set_page_config(layout='wide', page_title="Virtual Tutor")
-    st.markdown("""
-        <svg width="600" height="100">
-            <text x="50%" y="50%" font-family="San serif" font-size="42px" fill="Black" text-anchor="middle" stroke="white"
-             stroke-width="0.3" stroke-linejoin="round">Virtual Tutor - CHAT
-            </text>
-        </svg>
-    """, unsafe_allow_html=True)
-
-    if "messages" not in st.session_state:
-        st.session_state.messages = [
-            {"role": "assistant", "content": "Hi! How may I assist you today?"}
-        ]
-
-    st.markdown("""
-        <style> 
-        .stChatInputContainer > div {
-        background-color: #000000;
-        }
-        </style>
-        """, unsafe_allow_html=True)
-
-    for message in st.session_state.messages: # Display the prior chat messages
-        with st.chat_message(message["role"]):
-            st.write(message["content"])
-
-    for i, msg in enumerate(memory_storage.messages):
-        name = "user" if i % 2 == 0 else "assistant"
-        st.chat_message(name).markdown(msg.content)
-
-    if user_input := st.chat_input("User Input"):
-        with st.chat_message("user"):
-            st.markdown(user_input)
-
-        with st.spinner("Generating Response..."):
-            with st.chat_message("assistant"):
-                response = get_answer(user_input)
-                answer = response['result']
-                st.markdown(answer)
-                message = {"role": "assistant", "content": answer}
-                message_u = {"role": "user", "content": user_input}
-                st.session_state.messages.append(message_u)
-                st.session_state.messages.append(message)
-                display_images(user_input)
-                display_videos_streamlit(user_input)
-
-
 def display_images(image_collection, query_text, max_distance=None, debug=False):
     """
     Display images in a Streamlit app based on a query.
@@ -138,4 +89,159 @@ def display_videos_streamlit(video_collection, query_text, max_distance=None, ma
             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, video_collection, user_query):
+    # Get frame candidates from the video collection
+    frame_candidates = frame_uris(video_collection, user_query, max_distance=1.55)
+    
+    # Get image candidates from the image collection
+    image_candidates = image_uris(image_collection, user_query, max_distance=1.5)
+
+    # Initialize the inputs dictionary with just the query
+    inputs = {"query": user_query}
+
+    # Add the frame if found
+    frame = frame_candidates[0] if frame_candidates else ""
+    inputs["frame"] = frame
+
+    # If image candidates exist, process the first image
+    if image_candidates:
+        image = image_candidates[0]
+        with PILImage.open(image) as img:
+            img = img.resize((img.width // 6, img.height // 6))  # Resize the image
+            img = img.convert("L")  # Convert to grayscale
+            with io.BytesIO() as output:
+                img.save(output, format="JPEG", quality=60)  # Save as JPEG with compression
+                compressed_image_data = output.getvalue()
+
+        # Encode the compressed image as base64
+        inputs["image_data_1"] = base64.b64encode(compressed_image_data).decode('utf-8')
+    else:
+        inputs["image_data_1"] = ""
+
+    return inputs
+
+
+def home():
+    st.header("Welcome")
+    #st.set_page_config(layout='wide', page_title="Virtual Tutor")
+    st.markdown("""
+        <svg width="600" height="100">
+            <text x="50%" y="50%" font-family="San serif" font-size="42px" fill="Black" text-anchor="middle" stroke="white"
+             stroke-width="0.3" stroke-linejoin="round">Virtual Tutor - CHAT
+            </text>
+        </svg>
+    """, unsafe_allow_html=True)
+
+    if "messages" not in st.session_state:
+        st.session_state.messages = [
+            {"role": "assistant", "content": "Hi! How may I assist you today?"}
+        ]
+
+    st.markdown("""
+        <style> 
+        .stChatInputContainer > div {
+        background-color: #000000;
+        }
+        </style>
+        """, unsafe_allow_html=True)
+
+    for message in st.session_state.messages: # Display the prior chat messages
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
+
+    for i, msg in enumerate(memory_storage.messages):
+        name = "user" if i % 2 == 0 else "assistant"
+        st.chat_message(name).markdown(msg.content)
+
+    if user_input := st.chat_input("User Input"):
+        with st.chat_message("user"):
+            st.markdown(user_input)
+
+        with st.spinner("Generating Response..."):
+            with st.chat_message("assistant"):
+                response = get_answer(user_input)
+                answer = response['result']
+                st.markdown(answer)
+                
+                message = {"role": "assistant", "content": answer}
+                message_u = {"role": "user", "content": user_input}
+                st.session_state.messages.append(message_u)
+                st.session_state.messages.append(message)
+                inputs = format_prompt_inputs(image_collection,video_collection, user_input)
+                st.markdown("### Images")
+                display_images(image_collection, query, max_distance=1.55, debug=False)
+                st.markdown("### Videos")
+                frame = inputs["frame"]
+                if frame:
+                    directory_name = frame.split('/')[1] 
+                    video_path = f"videos_flattened/{directory_name}.mp4"
+                    if os.path.exists(video_path):
+                        st.video(video_path)
+                    else:
+                        st.write("No related videos found.")
+