Update app.py

app.py

@@ -1,12 +1,53 @@
 # app.py
 import streamlit as st
 import tensorflow as tf
-from tensorflow.keras.models import load_model
+from tensorflow import keras
 import numpy as np
 from PIL import Image
+import matplotlib.pyplot as plt
+from sklearn.model_selection import train_test_split
 
-# Load the pre-trained model
-model = load_model("cifar10_cnn_model.h5")
+# Load CIFAR-10 dataset
+(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
+
+# Normalize pixel values to [0, 1]
+x_train, x_test = x_train / 255.0, x_test / 255.0
+
+# Split training data into training and validation sets
+x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.2, random_state=42)
+
+# Define a simple CNN model
+def create_model():
+    model = keras.models.Sequential([
+        keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
+        keras.layers.MaxPooling2D((2, 2)),
+        keras.layers.Conv2D(64, (3, 3), activation='relu'),
+        keras.layers.MaxPooling2D((2, 2)),
+        keras.layers.Conv2D(128, (3, 3), activation='relu'),
+        keras.layers.Flatten(),
+        keras.layers.Dense(128, activation='relu'),
+        keras.layers.Dense(10, activation='softmax')
+    ])
+    return model
+
+# Check if the model is already saved
+import os
+if not os.path.exists("cifar10_cnn_model.h5"):
+    # Create and compile the model
+    model = create_model()
+    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+
+    # Train the model
+    st.write("Training the model...")
+    history = model.fit(x_train, y_train, epochs=10, validation_data=(x_val, y_val))  # Reduced epochs for quick testing
+
+    # Save the model
+    model.save("cifar10_cnn_model.h5")
+    st.write("Model saved as 'cifar10_cnn_model.h5'")
+else:
+    # Load the pre-trained model
+    st.write("Loading pre-trained model...")
+    model = keras.models.load_model("cifar10_cnn_model.h5")
 
 # Class names for CIFAR-10 dataset
 class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']