Update app.py

app.py

@@ -1,77 +1,80 @@
-import streamlit as st
+import streamlit as st 
 import tensorflow as tf
 import numpy as np
 import os
 import pandas as pd
 import altair as alt
 from PIL import Image
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
 import zipfile
 import io
 import shutil
 
-# 🎨 App Styling
-st.set_page_config(page_title="Cat vs Dog Classifier", page_icon="🐶🐱", layout="wide")
-st.markdown("""
-    <style>
-        .stApp { background-color: #f4f4f4; }
-        .stTitle { color: #333; text-align: center; }
-        .stMarkdown h2 { color: #555; }
-    </style>
-""", unsafe_allow_html=True)
-
-# 🎯 App Title
+# 🎨 App Title
 st.title("🐶🐱 Cat vs Dog Classifier")
 
 # 📖 About the App
-st.markdown(
+st.write(
     """
     ## About This App
-    This machine learning application classifies images into two categories: **Cats 🐱** and **Dogs 🐶**.
-    It is powered by **Streamlit** for an interactive user experience and **TensorFlow** for image classification.
-    
-    ### 🔥 Features:
-    - 📁 **Dataset Overview**: Displays dataset statistics and image distribution.
-    - 📊 **Model Performance**: Shows accuracy and loss metrics.
-    - 📷 **Image Classification**: Upload an image, and the model predicts the category with confidence.
-    - 📦 **Download Test Folder**: Option to download a ZIP of test images.
-    """,
-    unsafe_allow_html=True
+    This is a machine learning application that classifies images into two categories:
+    **Cats 🐱** and **Dogs 🐶**. The model is trained using a deep learning architecture
+    called Convolutional Neural Networks (CNNs) and is able to distinguish between images 
+    of cats and dogs with high accuracy.
+    ### Features:
+    - **Dataset Overview**: View the number of images in the dataset, categorized by "Cats" and "Dogs".
+    - **Model Evaluation**: Check the model's performance on the validation set, including accuracy and loss.
+    - **Image Classification**: Upload an image, and the model will predict whether it's a cat or a dog, along with the confidence level.
+    - **Download Test Folder**: Download a ZIP file containing the test images.
+    The app is powered by **Streamlit** for an interactive user interface and **TensorFlow** for image classification.
+    """
 )
 
-# ✅ Set Dataset Path
-BASE_DIR = "dataset"
-ZIP_PATH = "dataset.zip"
+# ✅ Detect Environment & Set Dataset Path
+BASE_DIR = "dataset"  # In Hugging Face Spaces, the dataset folder should be at the root of the Space
+ZIP_PATH = "dataset.zip"  # If dataset is uploaded as a ZIP (make sure it's in the same directory as app.py)
+
 TRAIN_DIR = os.path.join(BASE_DIR, "train")
 TEST_DIR = os.path.join(BASE_DIR, "test")
-MODEL_PATH = "cats_dogs_model.h5"
-IMG_SIZE = (150, 150)
-BATCH_SIZE = 32
 
-# ✅ Extract Dataset if Needed
-if ZIP_PATH and os.path.exists(ZIP_PATH) and not os.path.exists(BASE_DIR):
-    with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
-        zip_ref.extractall(BASE_DIR)
+# ✅ Extract Dataset if Needed (Hugging Face)
+if ZIP_PATH and os.path.exists(ZIP_PATH):
+    if not os.path.exists(BASE_DIR):  # Avoid re-extracting
+        with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
+            zip_ref.extractall(BASE_DIR)  # Extract into dataset folder
     st.success("✅ Dataset extracted!")
 
-# 📌 Verify Dataset
-cat_dir, dog_dir = os.path.join(TRAIN_DIR, "cats"), os.path.join(TRAIN_DIR, "dogs")
+# 📌 Check if dataset exists
+if not os.path.exists(TRAIN_DIR):
+    st.error(f"❌ Dataset folder 'train' not found at {TRAIN_DIR}. Please upload the dataset.")
+    st.stop()
+
+# 📌 Verify Cats & Dogs Folders
+cat_dir = os.path.join(TRAIN_DIR, "cats")
+dog_dir = os.path.join(TRAIN_DIR, "dogs")
+
 if not os.path.exists(cat_dir) or not os.path.exists(dog_dir):
-    st.error("❌ Missing 'cats' or 'dogs' folders in 'train'. Check your dataset.")
+    st.error("❌ Missing 'cats' or 'dogs' folders inside 'train'. Please check your dataset.")
     st.stop()
 
+# 📌 Constants
+IMG_SIZE = (150, 150)
+BATCH_SIZE = 32
+MODEL_PATH = "cats_dogs_model.h5"  # Ensure the model is uploaded to Hugging Face Space
+
 # 🎯 Load Model
 if os.path.exists(MODEL_PATH):
     model = tf.keras.models.load_model(MODEL_PATH)
 else:
-    st.error("⚠ No trained model found. Upload 'cats_dogs_model.h5'.")
+    st.error("⚠ No trained model found. Please upload 'cats_dogs_model.h5' to your Hugging Face repository.")
     st.stop()
 
 # 📷 Image Preprocessing
-@st.cache_resource
 def preprocess_image(image):
     image = image.convert('RGB').resize(IMG_SIZE)
     img_array = np.array(image, dtype=np.float32) / 255.0
-    return np.expand_dims(img_array, axis=0)
+    img_array = np.expand_dims(img_array, axis=0)
+    return img_array
 
 # 🐶🐱 Classify Image
 def classify_image(image):
@@ -83,63 +86,87 @@ def classify_image(image):
 
 # 📊 Model Evaluation
 def evaluate_model():
-    datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
-    test_data = datagen.flow_from_directory(TRAIN_DIR, target_size=IMG_SIZE, batch_size=BATCH_SIZE, class_mode='binary', subset='validation')
+    datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
+    
+    test_data = datagen.flow_from_directory(
+        TRAIN_DIR,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode='binary',
+        subset='validation'
+    )
+
     loss, accuracy = model.evaluate(test_data, verbose=0)
+
     return accuracy, loss
 
 # 🎛 Streamlit Tabs
 tab1, tab2, tab3 = st.tabs(["📁 Dataset Preview", "📊 Model Performance", "🐶🐱 Image Classification"])
 
-# 📁 Tab 1: Dataset Overview
+# 📁 Tab 1: Dataset Preview
 with tab1:
     st.write("### Dataset Overview")
+
     dataset_info = {
         "Total Images": len(os.listdir(cat_dir)) + len(os.listdir(dog_dir)),
         "Cat Images": len(os.listdir(cat_dir)),
         "Dog Images": len(os.listdir(dog_dir))
     }
-    df_info = pd.DataFrame(dataset_info.items(), columns=["Category", "Count"])
-    st.dataframe(df_info, width=500)
     
+    df_info = pd.DataFrame(list(dataset_info.items()), columns=["Category", "Count"])
+    st.dataframe(df_info)
+
+    # Visualization
     st.write("### Image Distribution")
-    chart = alt.Chart(df_info).mark_bar().encode(x="Category", y="Count", color="Category")
+    chart = alt.Chart(df_info).mark_bar().encode(
+        x="Category", y="Count", color="Category"
+    )
     st.altair_chart(chart, use_container_width=True)
 
 # 📊 Tab 2: Model Performance
 with tab2:
     st.write("### Model Evaluation")
+    
     accuracy, loss = evaluate_model()
-    st.metric(label="Validation Accuracy", value=f"{accuracy*100:.2f}%")
-    st.metric(label="Validation Loss", value=f"{loss:.4f}")
+    
+    st.write(f"✅ **Validation Accuracy:** {accuracy*100:.2f}%")
+    st.write(f"✅ **Validation Loss:** {loss:.4f}")
 
 # 🐶🐱 Tab 3: Image Classification
 with tab3:
     st.write("### Upload an Image for Classification")
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
-    
+
     if uploaded_file:
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Image", use_container_width=True)
-        
+
         with st.spinner("Classifying..."):
             label, confidence = classify_image(image)
-        
-        st.success(f"Prediction: **{label}** with **{confidence*100:.2f}%** confidence.")
-    
+
+        st.subheader("Prediction:")
+        st.write(f"This is a **{label}** with **{confidence*100:.2f}%** confidence.")
+
     # **New Feature: Download the 'test' folder as a ZIP**
     def zip_folder(folder_path):
+        # Create an in-memory zip file
         zip_buffer = io.BytesIO()
         with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
-            for root, _, files in os.walk(folder_path):
+            for root, dirs, files in os.walk(folder_path):
                 for file in files:
                     zip_file.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
-        zip_buffer.seek(0)
+        zip_buffer.seek(0)  # Go to the beginning of the file
         return zip_buffer
-    
+
+    # Button to download 'test' folder
     if os.path.exists(TEST_DIR):
         st.write("### Download Test Folder")
         zip_buffer = zip_folder(TEST_DIR)
-        st.download_button("Download Test Folder as ZIP", data=zip_buffer, file_name="test_folder.zip", mime="application/zip")
+        st.download_button(
+            label="Download Test Folder as ZIP",
+            data=zip_buffer,
+            file_name="test_folder.zip",
+            mime="application/zip"
+        )
     else:
-        st.warning("❌ Test folder not found.")
+        st.warning(f"❌ Test folder not found at {TEST_DIR}")