Update user.py

user.py

@@ -62,4 +62,80 @@ def home():
                 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.
+    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])
+
+    # Display images side by side, 3 images per row
+    cols = st.columns(3)  # Create 3 columns for the layout
+
+    for i, (uri, distance) in enumerate(sorted_results):
+        if max_distance is None or distance <= max_distance:
+            try:
+                img = PILImage.open(uri)
+                with cols[i % 3]:  # Use modulo to cycle through columns
+                    st.image(img, use_container_width = True)
+            except Exception as e:
+                st.error(f"Error loading image: {e}")
+
+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)")
+
+