File size: 2,811 Bytes
5855c83 084e7b7 5855c83 084e7b7 5855c83 084e7b7 6e52efa 084e7b7 5855c83 9290447 5855c83 79fba81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
# app.py
import streamlit as st
import tensorflow as tf
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 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=40, 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']
# Streamlit app title
st.title("Image Detection System")
# Upload image
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Display the uploaded image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_column_width=True)
# Preprocess the image
image = image.resize((32, 32)) # Resize to match CIFAR-10 input size
image = np.array(image) / 255.0 # Normalize pixel values
image = np.expand_dims(image, axis=0) # Add batch dimension
# Make prediction
predictions = model.predict(image)
predicted_class = np.argmax(predictions)
confidence = np.max(predictions) * 100
# Display results
st.write(f"**Prediction:** {class_names[predicted_class]}")
st.write(f"**Confidence:** {confidence:.2f}%")
model.save("cifar10_cnn_model.keras") |