CatDog-Detector / app.py
UXChaitanya's picture
Update app.py
1b0ce2e verified
raw
history blame
1.85 kB
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)