import streamlit as st
from PIL import Image
import numpy as np
import keras

# Load pre-trained model
model = keras.models.load_model('./image_classification_model.keras')
image_size = (180, 180)

# Function to make prediction
def predict(image):
    image_size = (180, 180)
    img = keras.utils.load_img(image, target_size=image_size)

    img_array = keras.utils.img_to_array(img)
    img_array = np.expand_dims(img_array, 0)  # Create batch axis

    predictions = model.predict(img_array)
    score = float(keras.activations.sigmoid(predictions[0][0]))
    return score

# Streamlit app
def main():
    st.title("Image Classification from Scratch")
    st.write("Upload an image to predict whether the image contains a cat or a dog.")

    uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])

    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        if st.button('Predict'):
            score = predict(uploaded_image)
            if (1 - score) > score:
                st.write('Prediction Result: {:.2f}% Cat'.format(100 * (1 - score)))
            else:
                st.write('Prediction Result: {:.2f}% Dog'.format(100 * score))

if __name__ == '__main__':
    main()