Update app.py

app.py

@@ -8,75 +8,97 @@ from PIL import Image
 from tensorflow.keras.preprocessing.image import ImageDataGenerator
 import zipfile
 import io
+import shutil
 
 # 🎨 App Title
 st.title("🐶🐱 Cat vs Dog Classifier")
 
 # 📖 About the App
-st.write("""
+st.write(
+    """
     ## About This App
     This is a machine learning application that classifies images into two categories:
-    **Cats 🐱** and **Dogs 🐶**. The model is trained using Convolutional Neural Networks (CNNs)
-    and distinguishes between images with high accuracy.
-    
+    **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 image counts for "Cats" and "Dogs".
-    - **Model Evaluation**: Check accuracy and loss on the validation set.
-    - **Image Classification**: Upload an image to classify it.
-    - **Download Test Folder**: Get test images as a ZIP file.
-""")
-
-# ✅ Set Dataset Path
-BASE_DIR = "dataset"
-ZIP_PATH = "dataset.zip"
+    - **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.
+    """
+)
+
+# ✅ 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!")
 
-# 📌 Check dataset existence
+# 📌 Check if dataset exists
 if not os.path.exists(TRAIN_DIR):
-    st.error(f"❌ Dataset folder 'train' not found at {TRAIN_DIR}.")
+    st.error(f"❌ Dataset folder 'train' not found at {TRAIN_DIR}. Please upload the dataset.")
     st.stop()
 
-cat_dir, dog_dir = os.path.join(TRAIN_DIR, "cats"), os.path.join(TRAIN_DIR, "dogs")
+# 📌 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 inside 'train'.")
+    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
 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):
-    prediction = model.predict(preprocess_image(image))[0][0]
+    processed_img = preprocess_image(image)
+    prediction = model.predict(processed_img)[0][0]
     label = "Dog 🐶" if prediction > 0.5 else "Cat 🐱"
-    confidence = max(prediction, 1 - prediction)
+    confidence = prediction if label == "Dog 🐶" else 1 - prediction
     return label, confidence
 
 # 📊 Model Evaluation
 def evaluate_model():
-    test_data = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2).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'
     )
-    return model.evaluate(test_data, verbose=0)
+
+    loss, accuracy = model.evaluate(test_data, verbose=0)
+
+    return accuracy, loss
 
 # 🎛 Streamlit Tabs
 tab1, tab2, tab3 = st.tabs(["📁 Dataset Preview", "📊 Model Performance", "🐶🐱 Image Classification"])
@@ -84,16 +106,29 @@ tab1, tab2, tab3 = st.tabs(["📁 Dataset Preview", "📊 Model Performance", "
 # 📁 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"])
+
+    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(list(dataset_info.items()), columns=["Category", "Count"])
     st.dataframe(df_info)
-    st.altair_chart(alt.Chart(df_info).mark_bar().encode(x="Category", y="Count", color="Category"), use_container_width=True)
+
+    # Visualization
+    st.write("### Image Distribution")
+    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.write(f"✅ **Validation Accuracy:** {accuracy*100:.2f}%")
     st.write(f"✅ **Validation Loss:** {loss:.4f}")
 
@@ -101,27 +136,37 @@ with tab2:
 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.subheader("Prediction:")
         st.write(f"This is a **{label}** with **{confidence*100:.2f}%** confidence.")
 
-    # **Download Test Folder as ZIP**
+    # **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(f"❌ Test folder not found at {TEST_DIR}")
+        st.warning(f"❌ Test folder not found at {TEST_DIR}")