nihal177's picture
Update app.py
56ab56f verified
# 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()