sync with remote

app.py

@@ -2,6 +2,8 @@ import streamlit as st
 from datasets import load_dataset
 import numpy as np
 import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
 
 # Load dataset from Hugging Face
 repo_id = "louiecerv/cats_dogs_dataset"
@@ -10,53 +12,50 @@ dataset = load_dataset(repo_id)
 # Select the split (train set by default)
 split = "train"
 data = dataset[split]
+label_names = dataset[split].features["label"].names  # Get class names upfront
+
+def display_images(images, labels, label_names, cols=5):
+    """Helper function to display images in a grid."""
+    rows = (len(images) + cols - 1) // cols  # Calculate rows dynamically
+    fig, axes = plt.subplots(rows, cols, figsize=(2*cols, 2*rows)) #Adjust figure size dynamically
+    axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes] #Handle single plot case
+    for i, ax in enumerate(axes):
+        if i < len(images):
+            ax.imshow(images[i])
+            ax.set_title(label_names[labels[i]])
+            ax.axis("off")
+        else:
+            ax.axis("off")  # Turn off unused subplots
+
+    st.pyplot(fig)
+
 
 def main():
     # Streamlit app
     st.title("Image Dataset Explorer")
-    st.subheader(f"Displaying the first 25 images from the {split} set")
-
-    # Extract images and labels
-    images = np.array(data["image"][:25])  # First 25 images
-    labels = np.array(data["label"][:25])
-    label_names = dataset[split].features["label"].names  # Get class names
+    st.subheader(f"Displaying images from the {split} set")
 
-    # Display 5x5 grid of images
-    fig, axes = plt.subplots(5, 5, figsize=(10, 10))
-    fig.subplots_adjust(hspace=0.5)
-
-    for i, ax in enumerate(axes.flat):
-        ax.imshow(images[i])
-        ax.set_title(label_names[labels[i]])
-        ax.axis("off")
-
-    st.pyplot(fig)
+    #Initial Images
+    if st.button("Show First 25 Images"):
+        with st.spinner("Loading images..."):
+            images = np.array(data["image"][:25])
+            labels = np.array(data["label"][:25])
+            display_images(images, labels, label_names)
 
     # Additional Exploration Options
     st.sidebar.title("Explore the Dataset")
 
     # Random Image Viewer
     if st.sidebar.button("Show Random Images"):
-        rand_indices = np.random.choice(len(data), 25, replace=False)
-        rand_images = np.array(data["image"])[rand_indices]
-        rand_labels = np.array(data["label"])[rand_indices]
-
-        fig, axes = plt.subplots(5, 5, figsize=(10, 10))
-        fig.subplots_adjust(hspace=0.5)
-
-        for i, ax in enumerate(axes.flat):
-            ax.imshow(rand_images[i])
-            ax.set_title(label_names[rand_labels[i]])
-            ax.axis("off")
-
-        st.pyplot(fig)
+        with st.spinner("Loading images..."):
+            rand_indices = np.random.choice(len(data), 25, replace=False)
+            rand_images = np.array(data["image"])[rand_indices]
+            rand_labels = np.array(data["label"])[rand_indices]
+            display_images(rand_images, rand_labels, label_names)
 
     # Class Distribution
     if st.sidebar.button("Show Class Distribution"):
-        import pandas as pd
-        import seaborn as sns
-
-        label_counts = pd.Series(labels).value_counts().sort_index()
+        label_counts = pd.Series(data["label"]).value_counts().sort_index() #Use entire dataset for distribution
         label_names_map = {i: name for i, name in enumerate(label_names)}
 
         fig, ax = plt.subplots(figsize=(8, 4))
@@ -64,27 +63,20 @@ def main():
         ax.set_title("Class Distribution")
         ax.set_ylabel("Count")
         ax.set_xlabel("Class")
-
         st.pyplot(fig)
 
     # Filter by class label
     selected_label = st.sidebar.selectbox("Filter by Label", label_names)
     if st.sidebar.button("Show Filtered Images"):
-        filtered_indices = [i for i, label in enumerate(labels) if label_names[label] == selected_label]
-        filtered_images = np.array(data["image"])[filtered_indices[:25]]
-        
-        fig, axes = plt.subplots(5, 5, figsize=(10, 10))
-        fig.subplots_adjust(hspace=0.5)
-
-        for i, ax in enumerate(axes.flat):
-            if i < len(filtered_images):
-                ax.imshow(filtered_images[i])
-                ax.set_title(selected_label)
-                ax.axis("off")
+        with st.spinner("Loading images..."):
+            filtered_indices = [i for i, label in enumerate(data["label"]) if label_names[label] == selected_label]
+            filtered_images = np.array(data["image"])[filtered_indices] #No need to slice if displaying all found
+            filtered_labels = np.array(data["label"])[filtered_indices]
+            if len(filtered_images) > 0: #check for no images found
+                display_images(filtered_images, filtered_labels, label_names)
             else:
-                ax.axis("off")
+                st.write(f"No images found for label: {selected_label}")
 
-        st.pyplot(fig)
 
 if __name__ == "__main__":
-    main()
+    main()