Update app.py

app.py

@@ -3,31 +3,42 @@ import tensorflow as tf
 from tensorflow.keras.preprocessing import image
 import numpy as np
 
-# Load the trained model
-model = tf.keras.models.load_model('Model_catsVSdogs.h5')
-
-# Define a function to make predictions
-def predict_image(img):
-    # Preprocess the image
-    img = img.resize((256,256))  # Resize the image to 224x224 pixels
-    img_array = image.img_to_array(img)  # Convert the image to an array
-    img_array = np.expand_dims(img_array, axis=0)  # Add a batch dimension
-    img_array = img_array / 255.0  # Normalize the image
-
-    # Make a prediction
+
+def load_models():
+    models = {}
+    models['SimpleNN_model'] = tf.keras.models.load_model("SimpleNN_model.h5")
+    models['VGG16'] = tf.keras.models.load_model("vgg16.h5")  
+    return models
+
+
+models = load_models()
+
+
+def predict_image(img, model_name):
+    model = models[model_name]
+    
+    if model_name == 'SimpleNN_model':
+        img = img.resize((256, 256))
+    elif model_name == 'VGG16':
+        img = img.resize((224, 224))
+
+    img_array = image.img_to_array(img)  
+    img_array = np.expand_dims(img_array, axis=0)  
+    img_array = img_array / 255.0  
+
+
     prediction = model.predict(img_array)
     if prediction[0] < 0.5:
         return "Cat"
     else:
         return "Dog"
 
-# Create the Gradio interface
 
 interface = gr.Interface(fn=predict_image, 
-                         inputs=gr.Image(type="pil"), 
+                         inputs=[gr.Image(type="pil"), gr.Dropdown(["SimpleNN_model", "VGG16"], label="Select Model")], 
                          outputs="text",
                          title="Cat and Dog Classifier",
-                         description="Upload an image of a cat or dog and the model will predict which one it is.")
+                         description="Upload an Image")
+
 
-# Launch the interface
-interface.launch(share=True)
+interface.launch()