nihal177's picture
Upload 5 files
5b3aad0 verified
raw
history blame
2.2 kB
# 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_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)
def load_models():
model_name = 'Model/mango_model.h5'
model = tf.keras.models.load_model(model_name)
return model
def load_labels():
with open('dataset_labels.txt', 'r') as file:
data = file.read().splitlines()
mango_dict = dict(enumerate(data, 1))
return mango_dict
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, categories, 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)
prediction_dict = dict(enumerate(prediction.flatten(), 1))
k = Counter(prediction_dict)
# Finding 3 highest values
high = k.most_common(3)
percentages = []
flowers = []
for i in high:
key, value = i
flowers.append(categories[key])
percentages.append(np.round(value*100, 2))
return flowers, percentages
def main():
st.title('Mango Ripeness Classifier 🥭')
model = load_models()
categories = load_labels()
image = load_image()
result = st.button('Run on image')
if result:
st.write('Calculating results...')
flowers, percentages = predict(model, categories, image)
st.text(flowers)
st.text(percentages)
if __name__ == '__main__':
main()