Vikranth's picture
Update app.py
3c9c18e
import tensorflow as tf
from keras.losses import SparseCategoricalCrossentropy
from keras.metrics import SparseCategoricalAccuracy
from PIL import Image
import numpy as np
from huggingface_hub import from_pretrained_keras
import gradio as gr
# prepare model
model = from_pretrained_keras("viola77data/recycling")
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5)
cls_loss = SparseCategoricalCrossentropy()
cls_acc = SparseCategoricalAccuracy()
model.compile(optimizer=optimizer, loss=cls_loss, metrics=[cls_acc])
# prepare the categories
categories = ['aluminium', 'batteries', 'cardboad',
'disposable plates', 'glass', 'hard plastic',
'paper', 'paper towel', 'polystyrene',
'soft plastics', 'takeaway cups']
dict_recycle = {
'aluminium': 'recycle',
'batteries': 'recycle',
'cardboad': 'recycle',
'disposable plates': 'dont recycle',
'glass': 'recycle',
'hard plastic': 'recycle',
'paper': 'recycle',
'paper towel': 'recycle',
'polystyrene': ' dont recycle',
'soft plastics': 'dont recycle',
'takeaway cups': 'dont recycle'
}
# prediction functions
def preprocess_image(im):
""" Pass in a numpy image an it returns a
TF Image"""
im = tf.cast(im, tf.float32) / 255.0
if len(im.shape) < 3:
im = tf.expand_dims(im, axis=-1) # add the channel dimension
im = tf.image.grayscale_to_rgb(im)
im = tf.image.resize(im, (224, 224))
im = tf.expand_dims(im, axis=0)
return im
def classify_image(input):
input_processed = preprocess_image(input)
preds = model.predict(input_processed)[0]
cls_preds = dict(zip(categories, map(float, preds)))
predicted_class = categories[np.argmax(preds)]
recycle_preds = dict_recycle[predicted_class]
return cls_preds
# Defining the Gradio Interface
# This is how the Demo will look like.
title = "Should I Recycle This?"
# description = """
# This app was created to help people recycle the right type of waste.
# You can use it at the comfort of your own home. Just take a picture of the waste material you want to know if
# its recyclible and upload it to this app and using Artificial Intelligence it will determine if you should
# throw the waste in the recycling bin or the normal bin.
# """
image = gr.Image(shape=(224,224))
label = gr.Label(num_top_classes=3, label='Prediction Material')
#recycle = gr.Textbox(label='Should you recycle?')
outputs = [label]
intf = gr.Interface(fn=classify_image, inputs=image, outputs=outputs, title = title,
cache_examples=False)
intf.launch(enable_queue=True)