from tensorflow.keras.models import load_model
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
import numpy as np
import gradio as gr
# Load your trained model
model = load_model("best_model_finetuned224.keras")
# Class labels
class_names = ['Cardboard', 'Glass', 'Metal', 'Paper', 'Plastic', 'Trash']
# Classify image and return top prediction with styled output
def classify_image(img):
try:
if img is None:
return "No image received. Please upload or capture an image.
"
img = img.resize((224, 224))
img_array = np.array(img, dtype=np.float32)
if img_array.shape != (224, 224, 3):
return f"Unexpected input shape: {img_array.shape}
"
img_array = preprocess_input(img_array)
img_array = np.expand_dims(img_array, axis=0)
prediction = model.predict(img_array)[0]
predicted_class_index = np.argmax(prediction)
predicted_class_name = class_names[predicted_class_index]
confidence = float(prediction[predicted_class_index])
# Return styled HTML result
return f"""
Prediction: {predicted_class_name}
Confidence: {confidence:.2%}
"""
except Exception as e:
return f"Error during prediction: {str(e)}
"
# Gradio interface with webcam and upload support
iface = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil", sources=["upload", "webcam"],), #
outputs=gr.HTML(),
title="♻️ Waste Material Classifier",
description="Upload or capture an image and this AI model will classify it into one of six recyclable waste categories.",
article="""
Powered by EfficientNetV2
Built with TensorFlow and Gradio
Developed by Saaiem Salaar
