Update app.py

app.py

@@ -1,57 +1,56 @@
+#denis_mnist_cnn_model_resnet50_v1.h5")  # Ensure you upload this file to Hugging Face Spaces
+
 import gradio as gr
 import tensorflow as tf
-from tensorflow.keras.applications.resnet import ResNet152, preprocess_input, decode_predictions
-from tensorflow.keras.preprocessing.image import img_to_array
-from PIL import Image
 import numpy as np
+from tensorflow.keras.applications.resnet50 import preprocess_input
+from tensorflow.keras.utils import load_img, img_to_array
+
+# Load your trained model
+model = tf.keras.models.load_model("denis_mnist_cnn_model_resnet50_v1.h5")  # Ensure you upload this file to Hugging Face Spaces
 
-# Load the pre-trained ResNet152 model
-MODEL_PATH = "resnet152-image-classifier.h5"  # Path to the saved model
-try:
-    model = tf.keras.models.load_model(MODEL_PATH)
-except Exception as e:
-    print(f"Error loading model: {e}")
-    exit()
-
-def decode_image_from_base64(base64_str):
-    # Decode the base64 string to bytes
-    image_data = base64.b64decode(base64_str)
-    # Convert the bytes into a PIL image
-    image = Image.open(BytesIO(image_data))
+# Define a function to preprocess the image
+def preprocess_image(image):
+    """
+    Preprocesses the uploaded image for prediction.
+    """
+    image = image.resize((128, 128))  # Resize to match the model input size
+    image = img_to_array(image)  # Convert PIL image to NumPy array
+    image = preprocess_input(image)  # Normalize for ResNet50
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
     return image
-    
-def predict_image(image):
+
+# Define the prediction function
+def predict(image):
     """
-    Process the uploaded image and return the top 3 predictions.
+    Accepts an image, preprocesses it, and returns the predicted label.
     """
-    try:
-        # Preprocess the image
-        image = image.resize((224, 224))  # ResNet152 expects 224x224 input
-        image_array = img_to_array(image)
-        image_array = preprocess_input(image_array)  # Normalize the image
-        image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
+    processed_image = preprocess_image(image)
+    predictions = model.predict(processed_image)
+    predicted_class = np.argmax(predictions, axis=-1)[0]  # Get the class index
+    confidence = np.max(predictions)  # Get confidence score
+    #return f"Predicted Class: {predicted_class}, Confidence: {confidence:.2f}"
+    return {"prediction": int(predicted_class)}
 
-        # Get predictions
-        predictions = model.predict(image_array)
-        decoded_predictions = decode_predictions(predictions, top=3)[0]
 
-        # Format predictions as a list of tuples (label, confidence)
-        results = [(label, float(confidence)) for _, label, confidence in decoded_predictions]
-        return dict(results)
+# Create a Gradio interface
+#interface = gr.Interface(fn=predict, inputs="image", outputs="json")
 
-    except Exception as e:
-        return {"Error": str(e)}
 
-# Create the Gradio interface
+# Create a Gradio interface
 interface = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil"),  # Accepts an image input
-    outputs=gr.Label(num_top_classes=3),  # Shows top 3 predictions with confidence
-    title="ResNet152 Image Classifier",
-    description="Upload an image, and the model will predict what's in the image.",
-    examples=["dog.jpg", "cat.jpg"],  # Example images for users to test
+    fn=predict,  # The prediction function
+    inputs=gr.Image(type="pil", label="Upload an Image"),  # Input: Image
+    outputs=gr.Textbox(label="Prediction"),  # Output: Textbox
+    title="MNIST ResNet50 Classifier",
+    description="Upload an image to classify it using the trained ResNet50 model.",
+    examples=[
+        ["example_images/example1.png"],  # Add paths to example images in your Hugging Face repository
+        ["example_images/example2.png"]
+    ],
 )
 
-# Launch the Gradio app
+
+# Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(share=True)