Spaces:
Sleeping
Sleeping
import streamlit as st | |
import tensorflow as tf | |
from keras.preprocessing.image import ImageDataGenerator | |
from keras.preprocessing import image | |
import numpy as np | |
# Check TensorFlow version | |
st.write("TensorFlow version: ", tf.__version__) | |
# Sidebar for file upload | |
st.sidebar.title("Upload Image") | |
uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg"]) | |
if uploaded_file is not None: | |
# Display the uploaded image | |
st.image(uploaded_file, caption='Uploaded Image.', use_column_width=True) | |
st.write("") | |
# Data preprocessing | |
train_datagen = ImageDataGenerator( | |
rescale=1./255, | |
shear_range=0.2, | |
zoom_range=0.2, | |
horizontal_flip=True | |
) | |
# Load the pre-trained model | |
cnn = tf.keras.models.Sequential([ | |
tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation='relu', input_shape=[64, 64, 3]), | |
tf.keras.layers.MaxPool2D(pool_size=2, strides=2), | |
tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation='relu'), | |
tf.keras.layers.MaxPool2D(pool_size=2, strides=2), | |
tf.keras.layers.Flatten(), | |
tf.keras.layers.Dense(units=128, activation='relu'), | |
tf.keras.layers.Dense(units=1, activation='sigmoid') | |
]) | |
# Compile the model | |
cnn.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy']) | |
# Load the trained weights | |
cnn.load_weights('cat_dog_classifier_weights.h5') | |
# Make prediction | |
test_image = image.load_img(uploaded_file, target_size=(64, 64)) | |
test_image = image.img_to_array(test_image) | |
test_image = np.expand_dims(test_image, axis=0) | |
result = cnn.predict(test_image) | |
# Class indices | |
class_indices = {0: 'cat', 1: 'dog'} | |
# Predict | |
prediction = 'dog' if result[0][0] == 1 else 'cat' | |
st.write("Prediction:", prediction) | |