import streamlit as st
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image

st.title('Dog Classification')

# import the model

model = load_model('model_best2.hdf5')

# define the preprocessing function

def preprocess_image(image):
    image = image.resize((240, 240))  # resize the image to the desired dimensions
    image = image.convert("RGB")  # convert the image to RGB mode if needed
    image = np.array(image)  # convert the image to a NumPy array
    image = image / 255.0  # normalize the pixel values to the range of 0 to 1
    image = np.expand_dims(image, axis=0)  # add an extra dimension for batch size
    return image

# define the prediction function

def prediction(image):
    preprocessed_image = preprocess_image(image)
    classes = model.predict(preprocessed_image)
    predicted_class_index = np.argmax(classes)
    class_labels = ['Afghan', 'Bulldog', 'Chow']
    predicted_class = class_labels[predicted_class_index]
    return predicted_class

# file uploader

uploaded_file = st.file_uploader("Upload your Dog Picture.")

# result

if st.button('Predict'):
    if uploaded_file is None:
        st.write('Please upload your favorite dog to purchase picture first.')
    else:
        image = Image.open(uploaded_file)
        result = prediction(image)
        st.write('This Dog belongs to the {} class.'.format(result))