| from ast import mod | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| import numpy as np | |
| import tensorflow as tf | |
| # Load the pre-trained model from Hugging Face | |
| model_name = "iamsuman/household-waste-EfficientNetV2M" | |
| model_path = hf_hub_download(repo_id=model_name, filename="EfficientNetV2M.h5") | |
| model = tf.keras.models.load_model(model_path) | |
| # Define a manual mapping of label indices to human-readable labels | |
| index_to_label = { | |
| 0: "battery", | |
| 1: "biological", | |
| 2: "cardboard", | |
| 3: "clothes", | |
| 4: "glass", | |
| 5: "metal", | |
| 6: "paper", | |
| 7: "plastic", | |
| 8: "shoes", | |
| 9: "trash" | |
| } | |
| def classify_image(image): | |
| image = Image.fromarray(image.astype('uint8'), 'RGB') | |
| image = image.resize((400, 400)) | |
| # Convert the PIL Image to a format compatible with the model | |
| image = np.array(image) | |
| image = np.expand_dims(image, axis=0) # Add batch dimension | |
| image = tf.keras.applications.efficientnet.preprocess_input(image) | |
| # Make prediction | |
| preds = model.predict(image) | |
| # Retrieve the highest probability class label index | |
| predicted_class_idx = np.argmax(preds, axis=-1)[0] | |
| # Convert the index to the model's class label | |
| label = index_to_label.get(predicted_class_idx, "Unknown Label") | |
| return label.capitalize() | |
| # Create Gradio interface | |
| iface = gr.Interface(fn=classify_image, | |
| inputs=gr.Image(), # Accepts image of any size | |
| outputs=gr.Label(), | |
| title="Household Waste Classification with EfficientNetV2M", | |
| description="Upload an image of household waste, and the model will classify it.") | |
| # Launch the app | |
| iface.launch() | |