Update app.py

app.py

@@ -2,14 +2,20 @@ import streamlit as st
 import os
 import random
 
-# 🖼️ 1. Display two columns, each with a random image from the directory
+# 🖼️ 1. Display two columns, each with a random image from the current directory
 def display_images(image_dir):
-    """🖼️ Function 1: Displays two random images side by side for voting."""
+    """🖼️ Function 1: Displays two random images side by side for voting, filtered to only show image files."""
+    
     col1, col2 = st.columns(2)
 
-    # 🗂️ Load random images from directory
-    images = os.listdir(image_dir)
+    # 🗂️ Load and filter only image files from the directory
+    valid_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp')  # Common image file types
+    images = [f for f in os.listdir(image_dir) if f.lower().endswith(valid_extensions)]
     
+    if len(images) < 2:
+        st.error("Not enough images in the directory.")
+        return
+
     # 📸 Randomly select two images
     image1 = random.choice(images)
     image2 = random.choice(images)
@@ -57,7 +63,7 @@ def main():
     show_vote_history()
     
     # 🖼️ Display images for voting
-    image_dir = '.'  # Set your directory path with images
+    image_dir = '.'  # Current directory where the app is running
     display_images(image_dir)
 
 # 🚀 5. Run the app