updated the dataset handling

app.py

@@ -5,59 +5,62 @@ import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn as sns
 
-# Load dataset from Hugging Face
-repo_id = "louiecerv/cats_dogs_dataset"
-dataset = load_dataset(repo_id)
+# Load dataset from Hugging Face with caching
+def load_hf_dataset():
+    repo_id = "louiecerv/cats_dogs_dataset"
+    return load_dataset(repo_id)
 
-# Select the split (train set by default)
+dataset = load_hf_dataset()
 split = "train"
 data = dataset[split]
-label_names = dataset[split].features["label"].names  # Get class names upfront
+label_names = dataset[split].features["label"].names
+
+@st.cache_data
+def get_label_counts():
+    return pd.Series(data["label"]).value_counts().sort_index()
 
 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
+    """Display images in a grid."""
+    rows = (len(images) + cols - 1) // cols
+    fig, axes = plt.subplots(rows, cols, figsize=(2 * cols, 2 * rows))
+    axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
+    
     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
-
+            ax.axis("off")
+    
     st.pyplot(fig)
 
-
 def main():
-    # Streamlit app
     st.title("Image Dataset Explorer")
     st.subheader(f"Displaying images from the {split} set")
-
-    #Initial Images
+    
+    # Show 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])
+            images = [data[i]["image"] for i in range(25)]
+            labels = [data[i]["label"] for i in range(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"):
         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)
+            rand_indices = [int(i) for i in np.random.choice(len(data), 25, replace=False)]
+            images = [data[i]["image"] for i in rand_indices]
+            labels = [data[i]["label"] for i in rand_indices]
+            display_images(images, labels, label_names)
 
     # Class Distribution
     if st.sidebar.button("Show Class Distribution"):
-        label_counts = pd.Series(data["label"]).value_counts().sort_index() #Use entire dataset for distribution
+        label_counts = get_label_counts()
         label_names_map = {i: name for i, name in enumerate(label_names)}
-
+        
         fig, ax = plt.subplots(figsize=(8, 4))
         sns.barplot(x=[label_names_map[i] for i in label_counts.index], y=label_counts.values, ax=ax)
         ax.set_title("Class Distribution")
@@ -69,14 +72,13 @@ def main():
     selected_label = st.sidebar.selectbox("Filter by Label", label_names)
     if st.sidebar.button("Show Filtered Images"):
         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)
+            filtered_indices = [int(i) for i in np.where(np.array(data["label"]) == label_names.index(selected_label))[0]]
+            if len(filtered_indices) > 0:
+                images = [data[i]["image"] for i in filtered_indices]
+                labels = [data[i]["label"] for i in filtered_indices]
+                display_images(images, labels, label_names)
             else:
                 st.write(f"No images found for label: {selected_label}")
 
-
 if __name__ == "__main__":
-    main()
+    main()