Spaces:
Runtime error
Runtime error
import numpy as np | |
import matplotlib.pyplot as plt | |
from sklearn.metrics import classification_report | |
from tensorflow import keras | |
from tensorflow.keras import layers | |
from tensorflow.keras.datasets import mnist | |
import streamlit as st | |
import os | |
# Load the MNIST dataset | |
(train_images, train_labels), (test_images, test_labels) = mnist.load_data() | |
# Preprocess the data | |
train_images = train_images.reshape((60000, 28, 28, 1)).astype("float32") / 255 | |
test_images = test_images.reshape((10000, 28, 28, 1)).astype("float32") / 255 | |
# Convert labels to categorical format | |
train_labels = keras.utils.to_categorical(train_labels, 10) | |
test_labels = keras.utils.to_categorical(test_labels, 10) | |
# Define the CNN model | |
def create_model(): | |
model = keras.Sequential([ | |
layers.Conv2D(32, (3, 3), activation="relu", input_shape=(28, 28, 1)), | |
layers.MaxPooling2D((2, 2)), | |
layers.Conv2D(64, (3, 3), activation="relu"), | |
layers.MaxPooling2D((2, 2)), | |
layers.Conv2D(64, (3, 3), activation="relu"), | |
layers.Flatten(), | |
layers.Dense(64, activation="relu"), | |
layers.Dense(10, activation="softmax") | |
]) | |
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"]) | |
return model | |
# Streamlit UI | |
st.title("CNN for MNIST Classification") | |
# Check if model is saved | |
model_path = "mnist_cnn_model.h5" | |
# Custom callback for logging | |
class StreamlitLogger(keras.callbacks.Callback): | |
def on_epoch_end(self, epoch, logs=None): | |
if logs is none: | |
logs = {} | |
st.write(f"Epoch {epoch + 1}:") | |
st.write(f" Train Loss: {logs.get('loss'):.4f}") | |
st.write(f" Train Accuracy: {logs.get('accuracy'):.4f}") | |
st.write(f" Val Loss: {logs.get('val_loss'):.4f}") | |
st.write(f" Val Accuracy: {logs.get('val_accuracy'):.4f}") | |
if st.button("Train Model"): | |
model = create_model() | |
logger = StreamlitLogger() | |
with st.spinner("Training..."): | |
history = model.fit(train_images, train_labels, validation_data=(test_images, test_labels), epochs=10, batch_size=64, callbacks=[logger]) | |
# Save the model | |
model.save(model_path) | |
# Plot training loss and accuracy | |
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) | |
ax1.plot(history.history["loss"], label="Train Loss") | |
ax1.plot(history.history["val_loss"], label="Val Loss") | |
ax1.set_title("Training and Validation Loss") | |
ax1.set_xlabel("Epoch") | |
ax1.set_ylabel("Loss") | |
ax2.plot(history.history["accuracy"], label="Train Accuracy") | |
ax2.plot(history.history["val_accuracy"], label="Val Accuracy") | |
ax2.set_title("Training and Validation Accuracy") | |
ax2.set_xlabel("Epoch") | |
ax2.set_ylabel("Accuracy") | |
ax1.legend() | |
ax2.legend() | |
st.pyplot(fig) | |
# Evaluate the model on test data | |
test_preds = np.argmax(model.predict(test_images), axis=1) | |
true_labels = np.argmax(test_labels, axis=1) | |
st.session_state['true_labels'] = true_labels | |
report = classification_report(true_labels, test_preds, digits=4) | |
st.text("Classification Report:") | |
st.text(report) | |
index = st.number_input("Enter an index (0-9999) to test:", min_value=0, max_value=9999, step=1) | |
def test_index_prediction(index): | |
image = test_images[index].reshape(28, 28) | |
st.image(image, caption=f"True Label: {st.session_state['true_labels'][index]}", use_column_width=True) | |
# Reload the model | |
if not os.path.exists(model_path): | |
st.error("Train the model first.") | |
return | |
model = keras.models.load_model(model_path) | |
prediction = model.predict(test_images[index].reshape(1, 28, 28, 1)) | |
predicted_class = np.argmax(prediction) | |
st.write(f"Predicted Class: {predicted_class}") | |
if st.button("Test Index"): | |
test_index_prediction(index) | |