Update app.py

app.py

@@ -8,97 +8,75 @@ 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 a deep learning architecture
-    called Convolutional Neural Networks (CNNs) and is able to distinguish between images 
-    of cats and dogs with high accuracy.
+    **Cats 🐱** and **Dogs 🐶**. The model is trained using Convolutional Neural Networks (CNNs)
+    and distinguishes between images 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.
-    """
-)
-
-# ✅ 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)
-
+    - **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"
 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 (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
+# ✅ 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)
     st.success("✅ Dataset extracted!")
 
-# 📌 Check if dataset exists
+# 📌 Check dataset existence
 if not os.path.exists(TRAIN_DIR):
-    st.error(f"❌ Dataset folder 'train' not found at {TRAIN_DIR}. Please upload the dataset.")
+    st.error(f"❌ Dataset folder 'train' not found at {TRAIN_DIR}.")
     st.stop()
 
-# 📌 Verify Cats & Dogs Folders
-cat_dir = os.path.join(TRAIN_DIR, "cats")
-dog_dir = os.path.join(TRAIN_DIR, "dogs")
-
+cat_dir, dog_dir = os.path.join(TRAIN_DIR, "cats"), 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'. Please check your dataset.")
+    st.error("❌ Missing 'cats' or 'dogs' folders inside 'train'.")
     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. Please upload 'cats_dogs_model.h5' to your Hugging Face repository.")
+    st.error("⚠ No trained model found. Upload 'cats_dogs_model.h5'.")
     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
-    img_array = np.expand_dims(img_array, axis=0)
-    return img_array
+    return np.expand_dims(img_array, axis=0)
 
 # 🐶🐱 Classify Image
 def classify_image(image):
-    processed_img = preprocess_image(image)
-    prediction = model.predict(processed_img)[0][0]
+    prediction = model.predict(preprocess_image(image))[0][0]
     label = "Dog 🐶" if prediction > 0.5 else "Cat 🐱"
-    confidence = prediction if label == "Dog 🐶" else 1 - prediction
+    confidence = max(prediction, 1 - prediction)
     return label, confidence
 
 # 📊 Model Evaluation
 def evaluate_model():
-    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'
+    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'
     )
-
-    loss, accuracy = model.evaluate(test_data, verbose=0)
-
-    return accuracy, loss
+    return model.evaluate(test_data, verbose=0)
 
 # 🎛 Streamlit Tabs
 tab1, tab2, tab3 = st.tabs(["📁 Dataset Preview", "📊 Model Performance", "🐶🐱 Image Classification"])
@@ -106,29 +84,16 @@ 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(list(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(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"
-    )
-    st.altair_chart(chart, use_container_width=True)
+    st.altair_chart(alt.Chart(df_info).mark_bar().encode(x="Category", y="Count", color="Category"), 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}")
 
@@ -136,37 +101,27 @@ 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.")
 
-    # **New Feature: Download the 'test' folder as a ZIP**
+    # **Download Test Folder as 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, dirs, files in os.walk(folder_path):
+            for root, _, 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)  # Go to the beginning of the file
+        zip_buffer.seek(0)
         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(
-            label="Download Test Folder as ZIP",
-            data=zip_buffer,
-            file_name="test_folder.zip",
-            mime="application/zip"
-        )
+        st.download_button("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}")