| from pathlib import Path
|
| import platform
|
| import pathlib
|
|
|
| import gradio as gr
|
| from fastai.vision.all import PILImage, load_learner
|
|
|
| if platform.system() == "Windows":
|
| pathlib.PosixPath = pathlib.WindowsPath
|
|
|
|
|
|
|
| MODEL_PATH = Path(__file__).parent / "models" / "waste_classifier_v2.pkl"
|
| learn = load_learner(MODEL_PATH)
|
|
|
|
|
| BIN_INFO = {
|
| "plastic": "Plastic Recycling",
|
| "glass": "Glass Recycling",
|
| "metal": "Metal Recycling",
|
| "paper": "Paper Recycling",
|
| "cardboard": "Paper/Cardboard Recycling",
|
| "organic": "Organic/Compost",
|
| "textile": "Textile Collection",
|
| "e-waste": "Electronic Waste Collection",
|
| }
|
|
|
|
|
| def predict(image):
|
| img = PILImage.create(image)
|
|
|
| pred, pred_idx, probs = learn.predict(img)
|
| confidence = float(probs[pred_idx])
|
|
|
| bin_name = BIN_INFO[str(pred)]
|
|
|
| return (
|
| f"## {pred.title()} ({confidence:.1%})",
|
| f"**Recommended disposal:** {bin_name}",
|
| )
|
|
|
|
|
| demo = gr.Interface(
|
| fn=predict,
|
| inputs=gr.Image(type="pil", label="Upload a waste image"),
|
| outputs=[
|
| gr.Markdown(label="Prediction"),
|
| gr.Markdown(label="Disposal Recommendation"),
|
| ],
|
| title="♻️ AI Trash Sorter",
|
| description="Upload an image of waste and the AI will identify it and recommend the correct trash bin.",
|
| )
|
|
|
| if __name__ == "__main__":
|
| demo.launch() |
