Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| from PIL import Image | |
| import numpy as np | |
| from tensorflow.keras.preprocessing.image import ImageDataGenerator, img_to_array | |
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout | |
| from tensorflow.keras.utils import to_categorical | |
| from tensorflow.keras.optimizers import Adam | |
| from sklearn.preprocessing import LabelEncoder | |
| from sklearn.model_selection import train_test_split | |
| # Page config | |
| st.set_page_config(page_title="Animal Classifier DL", layout="centered") | |
| st.title("π§ Animal Classifier (Deep Learning)") | |
| st.markdown("Train from folders like `Animals/cat`, `Animals/dog`, etc. and predict uploaded images.") | |
| IMAGE_SIZE = (128, 128) | |
| DATA_DIR = "Animals" # Folder with subfolders | |
| # π Load dataset and preprocess | |
| def load_dataset(): | |
| X = [] | |
| y = [] | |
| for label in os.listdir(DATA_DIR): | |
| folder = os.path.join(DATA_DIR, label) | |
| if os.path.isdir(folder): | |
| for file in os.listdir(folder): | |
| if file.lower().endswith(('.png', '.jpg', '.jpeg')): | |
| try: | |
| path = os.path.join(folder, file) | |
| img = Image.open(path).convert("RGB").resize(IMAGE_SIZE) | |
| arr = img_to_array(img) / 255.0 | |
| X.append(arr) | |
| y.append(label) | |
| except Exception as e: | |
| print(f"β οΈ Skipped {file}: {e}") | |
| return np.array(X), np.array(y) | |
| X, y = load_dataset() | |
| # Stop if no data | |
| if len(X) == 0: | |
| st.error("β No images found in 'Animals' folder.") | |
| st.stop() | |
| # Encode labels | |
| le = LabelEncoder() | |
| y_encoded = le.fit_transform(y) | |
| y_cat = to_categorical(y_encoded) | |
| # π§ Train/test split | |
| X_train, X_test, y_train, y_test = train_test_split(X, y_cat, test_size=0.2, random_state=42) | |
| # π§ Build model | |
| def build_model(): | |
| model = Sequential([ | |
| Conv2D(32, (3, 3), activation='relu', input_shape=(*IMAGE_SIZE, 3)), | |
| MaxPooling2D(2, 2), | |
| Conv2D(64, (3, 3), activation='relu'), | |
| MaxPooling2D(2, 2), | |
| Flatten(), | |
| Dense(128, activation='relu'), | |
| Dropout(0.3), | |
| Dense(len(np.unique(y_encoded)), activation='softmax') | |
| ]) | |
| model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy']) | |
| return model | |
| model = build_model() | |
| # π¨βπ« Train the model | |
| with st.spinner("Training model..."): | |
| model.fit(X_train, y_train, epochs=5, batch_size=32, validation_split=0.1, verbose=0) | |
| st.success("β Model trained!") | |
| # π€ Upload images for prediction | |
| uploaded_files = st.file_uploader("Upload animal images", type=["jpg", "jpeg", "png"], accept_multiple_files=True) | |
| if uploaded_files: | |
| st.markdown("### Predictions") | |
| cols = st.columns(len(uploaded_files[:3])) | |
| for i, file in enumerate(uploaded_files[:3]): | |
| with cols[i]: | |
| img = Image.open(file).convert("RGB") | |
| st.image(img, caption="Uploaded", use_container_width=True) | |
| img_resized = img.resize(IMAGE_SIZE) | |
| arr = img_to_array(img_resized) / 255.0 | |
| arr = np.expand_dims(arr, axis=0) | |
| pred = model.predict(arr, verbose=0)[0] | |
| top_idx = np.argmax(pred) | |
| label = le.inverse_transform([top_idx])[0] | |
| st.success(f"π Prediction: **{label.capitalize()}**") | |