feat: Add predict functionality

requirements.txt

@@ -1 +1,4 @@
-streamlit
+numpy==2.3.2
+pillow==11.3.0
+streamlit==1.48.1
+tensorflow==2.20.0

src/predictor.py

@@ -0,0 +1,50 @@
+import numpy as np
+from tensorflow.keras.applications.resnet50 import (
+    ResNet50,
+    decode_predictions,
+    preprocess_input,
+)
+from tensorflow.keras.preprocessing import image
+
+# Load the model outside the function to ensure it's loaded only once
+model = ResNet50(include_top=True, weights="imagenet")
+
+
+def predict_image(img):
+    """
+    Preprocesses an image and runs a pre-trained ResNet50 model to get a prediction.
+
+    Parameters
+    ----------
+    img : PIL.Image
+        The image object to classify.
+
+    Returns
+    -------
+    class_name, pred_probability : tuple(str, float)
+        The model's predicted class as a string and the corresponding confidence
+        score as a number.
+    """
+    # Resize the image to match model input dimensions (224, 224)
+    img = img.resize((224, 224))
+
+    # Convert Pillow image to np.array
+    x = image.img_to_array(img)
+
+    # Add an extra dimension for the batch size
+    x_batch = np.expand_dims(x, axis=0)
+
+    # Apply ResNet50-specific preprocessing
+    x_batch = preprocess_input(x_batch)
+
+    # Make predictions
+    predictions = model.predict(x_batch, verbose=0)
+
+    # Get predictions using model methods and decode predictions
+    top_pred = decode_predictions(predictions, top=1)[0][0]  # imagenet_id, label, score
+    _, class_name, pred_probability = top_pred
+
+    # Convert probability to float and round it
+    pred_probability = round(float(pred_probability), 4)
+
+    return class_name, pred_probability

src/streamlit_app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from PIL import Image
 
+from predictor import predict_image
+
 # 📌 PAGE SETUP
 st.set_page_config(page_title="Image Classifier App", page_icon="🤖", layout="centered")
 st.html("""
@@ -79,7 +81,11 @@ with st.container():
             if classify_button:
                 # Check if an image is selected before running prediction
                 if uploaded_image is not None:
-                    st.session_state["selected_image"] = uploaded_image
+                    # st.session_state["selected_image"] = uploaded_image
+                    # Use Image.open() to convert the UploadedFile object into a PIL.Image object
+                    st.session_state["selected_image"] = Image.open(uploaded_image)
+                    st.session_state["uploaded_file"] = uploaded_image
+
                 elif selected_example:
                     # Load the selected example image
                     try:
@@ -96,8 +102,6 @@ with st.container():
                         st.session_state["selected_image"],
                         caption="Image to be classified",
                     )
-                    st.markdown("---")
-                    st.subheader("Prediction")
 
                     # Call the prediction function and display results
                     with st.spinner("Analyzing image..."):
@@ -109,7 +113,7 @@ with st.container():
 
                             st.metric(
                                 label="Prediction",
-                                value=f"Prediction: {predicted_label.replace('_', ' ').title()}",
+                                value=f"{predicted_label.replace('_', ' ').title()}",
                                 delta=f"{predicted_score * 100:.2f}%",
                                 help="The predicted category and its confidence score.",
                                 delta_color="normal",