Update app.py

app.py

@@ -3,25 +3,26 @@ import gradio as gr
 import numpy as np
 import tensorflow as tf
 import cv2
+import tensorflow_hub as hub
 
 # Load class labels from the text file
 train_info = []
 with open('labelwithspace.txt', 'r') as file:
     train_info = [line.strip() for line in file.readlines()]
 
-# Construct a simple mock model using the Functional API
-def create_mock_model():
-    inputs = tf.keras.Input(shape=(224, 224, 3))
-    # A simple convolutional base (for demonstration purposes)
-    x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(inputs)
-    x = tf.keras.layers.MaxPooling2D((2, 2))(x)
-    x = tf.keras.layers.Flatten()(x)
-    outputs = tf.keras.layers.Dense(len(train_info), activation='softmax')(x)
-    model = tf.keras.Model(inputs, outputs)
+# Load your actual model from the .h5 file
+def load_real_model():
+    try:
+        # Register KerasLayer from TensorFlow Hub if used
+        custom_objects = {'KerasLayer': hub.KerasLayer}
+        model = tf.keras.models.load_model('bird_model4.h5', custom_objects=custom_objects)
+    except Exception as e:
+        print("Error loading the model:", e)
+        exit()
     return model
 
-# Initialize the mock model
-model = create_mock_model()
+# Initialize the real model
+model = load_real_model()
 
 # Function to preprocess the image and make predictions
 def predict_image(image):
@@ -29,16 +30,16 @@ def predict_image(image):
     img = cv2.resize(image, (224, 224))
     img = img / 255.0  # Normalize to [0, 1] range
 
-    # Make predictions using the mock model
+    # Make predictions using the loaded model
     predictions = model.predict(img[np.newaxis, ...])[0]
     top_classes = np.argsort(predictions)[-3:][::-1]  # Indices of top 3 predictions
     top_class = top_classes[0]  # Index of the highest probability class
     label = train_info[top_class]  # Get the corresponding label
     return label
 
-# Define the Gradio interface without the 'shape' argument
-input_image = gr.Image()  # Removed 'shape' argument
+# Define the Gradio interface
+input_image = gr.Image()
 output_label = gr.Label()
 
-# Launch the Gradio interface
-gr.Interface(fn=predict_image, inputs=input_image, outputs=output_label, capture_session=True).launch()
+# Launch the Gradio interface for Hugging Face deployment
+gr.Interface(fn=predict_image, inputs=input_image, outputs=output_label).launch()