SmartWaste / app.py
chaninder's picture
Update app.py
d44854d
raw
history blame
758 Bytes
import gradio as gr
import tensorflow as tf
from tensorflow import keras
# load pre-trained model
model_path = "/Users/chaninderrishi/Desktop/ML/projects/waste-sorting/models/prod3"
pre_trained_model = keras.models.load_model(model_path)
# classification labels
labels = ['compost', 'e-waste', 'recycle', 'trash']
def classify_image(input):
prediction = pre_trained_model.predict(input)
confidences = {labels[i]: float(prediction[i]) for i in range(4)}
return confidences
# create Gradio interface
iface = gr.Interface(fn=classify_image,
inputs=gr.Image(shape=(224, 224)),
outputs=gr.Label(num_top_classes=3),
#examples=["banana.jpg", "car.jpg"]
)
iface.launch(share=True)