# All imports
import streamlit as st
import tensorflow as tf
from tensorflow import keras
from PIL import Image
from tensorflow.keras.preprocessing import image
import io
from collections import Counter
import numpy as np

def load_models():
    model_name = 'Model/mango_new_model.h5'
    model = tf.keras.models.load_model(model_name) 
    return model
    
def load_image():
     uploaded_file = st.file_uploader(label='Pick an image to test')
     if uploaded_file is not None:
         image_data = uploaded_file.getvalue()
         st.image(image_data)
         img = Image.open(io.BytesIO(image_data))
         img = img.resize((224,224))
         return img
     else:
         return None
    
def predict(model, img):
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    prediction = [img_array]
    prediction_test = [1]
    test_ds = tf.data.Dataset.from_tensor_slices((prediction, prediction_test))
    test_ds = test_ds.cache().batch(32).prefetch(buffer_size = tf.data.experimental.AUTOTUNE)

    prediction = model.predict(test_ds)

    st.write(prediction)
    if prediction[0]>0.5:
        return 'ripe'
    else:
        return 'unripe'

def main():
    st.title('Mango Ripeness Classifier 🥭')
    model = load_models()
    image = load_image()
    result = st.button('Run on image')
    if result:
        st.write('Calculating results...')
        st.write(predict(model, image))

if __name__ == '__main__':
    main()