Update app.py

app.py

@@ -1,61 +1,55 @@
+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
-print(tf.__version__)
-
-# Data preprocessing
-train_datagen = ImageDataGenerator(
-    rescale=1./255,
-    shear_range=0.2,
-    zoom_range=0.2,
-    horizontal_flip=True
-)
-
-test_datagen = ImageDataGenerator(rescale=1./255)
-
-training_set = train_datagen.flow_from_directory(
-    'dataset/training_set',
-    target_size=(64, 64),
-    batch_size=32,
-    class_mode='binary'
-)
-
-test_set = test_datagen.flow_from_directory(
-    'dataset/test_set',
-    target_size=(64, 64),
-    batch_size=32,
-    class_mode='binary'
-)
-
-# Build CNN 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'])
-
-# Train the model
-cnn.fit(x=training_set, validation_data=test_set, epochs=25)
-
-# Single prediction
-test_image = image.load_img('dataset/single_prediction/cat_or_dog_1.jpg', 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 = training_set.class_indices
-
-# Predict
-prediction = 'dog' if result[0][0] == 1 else 'cat'
-print(prediction)
+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)