Spaces:
Sleeping
Sleeping
| # from __future__ import annotations | |
| import gradio as gr | |
| from huggingface_hub import from_pretrained_fastai | |
| import os | |
| from typing import Iterable | |
| # from gradio.themes.base import Base | |
| # from gradio.themes.utils import colors, fonts, sizes | |
| os.environ["HF_ENDPOINT"] = "https://huggingface.co" | |
| materials_model = from_pretrained_fastai("pyesonekyaw/recycletree_materials") | |
| paper_model = from_pretrained_fastai("pyesonekyaw/recycletree_paper") | |
| plastic_model = from_pretrained_fastai("pyesonekyaw/recycletree_plastic") | |
| metal_model = from_pretrained_fastai("pyesonekyaw/recycletree_metal") | |
| others_model = from_pretrained_fastai("pyesonekyaw/recycletree_others") | |
| glass_model = from_pretrained_fastai("pyesonekyaw/recycletree_glass") | |
| examples = ["Examples/1.jpg", "Examples/2.jpg", | |
| "Examples/3.jpg", "Examples/4.jpg", "Examples/5.jpg"] | |
| material_names = ['Glass', 'Metal', 'Others', 'Paper', 'Plastic'] | |
| plastic_names = ['CD Disk', 'Straw', 'Plastic Bag', 'Clothes Hanger', 'Plastic Container or Bottle', | |
| 'Disposable Cutlery', 'Plastic Packaging', 'Plastic Packaging With Foil', 'Styrofoam'] | |
| paper_names = ['Beverage Carton', 'Cardboard', 'Chopsticks', 'Disposables', 'Paper Bag', 'Paper Packaging', | |
| 'Paper Product', 'Receipt', 'Paper Roll', 'Paper Sheet', 'Tissue Box', 'Tissue Paper'] | |
| glass_names = ['Ceramic', 'Glassware', 'Lightbulb'] | |
| other_names = ['Battery', 'Electronic Waste', 'Stationery'] | |
| metal_names = ['Aerosol Can', 'Aluminium Foil or Tray', 'Metal Can or Container'] | |
| material_num_name_dict = { | |
| "metal": "Metal", | |
| "glass": "Glass", | |
| "paper": "Paper", | |
| "plastic": "Plastic", | |
| "others": "Others", | |
| } | |
| plastic_item_num_dict = { | |
| "CD Disk": ["Yes", "Nil"], | |
| "Straw": ["No, dispose as general waste","Nil"], | |
| "Plastic Bag": ["Yes, if they are not oxo- and bio- degradable bags", "Contaminated with food waste/liquid waste/other forms of waste. Please refer to RecycleHowAh for more information"], | |
| "Clothes Hanger": ["Yes", "Made up of more than one plastic, if unsure, just dispose as normal waste "], | |
| "Plastic Container or Bottle": ["Yes", "When they are not emptied or not rinsed "], | |
| "Disposable Cutlery": ["No, dispose as general waste", "Nil"], | |
| "Plastic Packaging": ["Yes, for things like bubble wrap and egg tray but no if directly enclosing food like cling wrap", "Contaminated with food contents "], | |
| "Plastic Packaging With Foil": ["No","Nil"], | |
| "Styrofoam": ["No, dispose as general waste","Nil"] | |
| } | |
| glass_item_num_dict = { | |
| "Ceramic": ["No, donate if can be reused", "Nil"], | |
| "Glassware": ["Yes","If there is liquid/solid residue inside the glassware "], | |
| "Lightbulb": ["Could be recycled at specific collection points which can be found on onemap.sg, under Lighting waste collection points", "Nil"] | |
| } | |
| metal_item_num_dict = { | |
| "Aerosol Can": ["Yes","If there are any remaining contents in the can"], | |
| "Aluminium Foil or Tray": ["Yes","If there is any residue "], | |
| "Metal Can or Container": ["Yes","If there is any residue "] | |
| } | |
| others_item_num_dict = { | |
| "Battery": ["Battery","No, rechargeable batteries can be recycled through specific collection points (e-waste collection)", "Nil"], | |
| "Electronic Waste": ["Electronic Waste","Can be recycled through specific collection points (e-waste collection). Please refer to RecycleHowAh for more information"], | |
| "Stationery": ["Stationery","No, donate if can be reused"] | |
| } | |
| paper_item_num_dict = { | |
| "Beverage Carton": ["Yes, rinsed and flattened","Nil"], | |
| "Cardboard": ["Yes","Remains of other materials such as tape, contaminated with other waste"], | |
| "Chopsticks": ["No, dispose as general waste ",], | |
| "Disposables": ["No, dispose as general waste ",], | |
| "Paper Bag": ["Yes","Contaminated with food waste or other waste "], | |
| "Paper Packaging": ["Yes","Made up of more than one material or contaminated with food waste"], | |
| "Paper Product": ["Yes","Contaminated with other waste"], | |
| "Receipt": ["Yes","Contaminated with other waste"], | |
| "Paper Roll": ["Yes","Contaminated with other waste"], | |
| "Paper Sheet": ["Yes","Contaminated with other waste "], | |
| "Tissue Box": ["Yes","Plastic liners not removed or contaminated with other waste "], | |
| "Tissue Paper": ["No, dispose as general waste","Nil"] | |
| } | |
| # class Seafoam(Base): | |
| # def __init__( | |
| # self, | |
| # *, | |
| # primary_hue: colors.Color | str = colors.emerald, | |
| # secondary_hue: colors.Color | str = colors.teal, | |
| # neutral_hue: colors.Color | str = colors.teal, | |
| # spacing_size: sizes.Size | str = sizes.spacing_md, | |
| # radius_size: sizes.Size | str = sizes.radius_md, | |
| # text_size: sizes.Size | str = sizes.text_lg, | |
| # font: fonts.Font | |
| # | str | |
| # | Iterable[fonts.Font | str] = ( | |
| # fonts.GoogleFont("Quicksand"), | |
| # "ui-sans-serif", | |
| # "sans-serif", | |
| # ), | |
| # font_mono: fonts.Font | |
| # | str | |
| # | Iterable[fonts.Font | str] = ( | |
| # fonts.GoogleFont("IBM Plex Mono"), | |
| # "ui-monospace", | |
| # "monospace", | |
| # ), | |
| # ): | |
| # super().__init__( | |
| # primary_hue=primary_hue, | |
| # secondary_hue=secondary_hue, | |
| # neutral_hue=neutral_hue, | |
| # spacing_size=spacing_size, | |
| # radius_size=radius_size, | |
| # text_size=text_size, | |
| # font=font, | |
| # font_mono=font_mono, | |
| # ) | |
| # super().set( | |
| # body_background_fill="linear-gradient(45deg, *primary_200, *primary_200 10px, *primary_50 10px, *primary_50 20px)", | |
| # body_background_fill_dark="linear-gradient(45deg, *primary_200, *primary_200 10px, *primary_50 10px, *primary_50 20px)", | |
| # stat_background_fill_dark="linear-gradient(to right, *primary_400, *primary_200)", | |
| # error_background_fill_dark=f"linear-gradient(to right, {colors.red.c100}, *background_fill_secondary)", | |
| # button_primary_background_fill="linear-gradient(90deg, *primary_300, *secondary_400)", | |
| # button_primary_background_fill_hover="linear-gradient(90deg, *primary_200, *secondary_300)", | |
| # button_primary_background_fill_hover_dark="linear-gradient(90deg, *primary_200, *secondary_300)", | |
| # button_primary_text_color="white", | |
| # button_primary_background_fill_dark="linear-gradient(90deg, *primary_300, *secondary_400)", | |
| # slider_color="*secondary_300", | |
| # slider_color_dark="*secondary_300", | |
| # block_title_text_weight="600", | |
| # block_border_width="3px", | |
| # block_shadow="*shadow_drop_lg", | |
| # button_shadow="*shadow_drop_lg", | |
| # button_large_padding="32px", | |
| # ) | |
| # for k in list(self.__dict__.keys()): | |
| # if '_dark' in k: | |
| # setattr(self,k,None) | |
| # seafoam = Seafoam() | |
| def predict_image(inp): | |
| """ | |
| Performs inference for a given input image and returns the prediction and CAM image. | |
| """ | |
| try: | |
| material_label, material_label_idx, material_probs = materials_model.predict(inp) | |
| material_preds = {name: prob for name, prob in zip(material_names, material_probs.tolist())} | |
| if material_label == 'paper': | |
| specific_label, specific_label_idx, specific_probs = paper_model.predict(inp) | |
| specific_preds = {name: prob for name, prob in zip(paper_names, specific_probs.tolist())} | |
| specific_label = paper_names[int(specific_label_idx)] | |
| recyclable_qn = paper_item_num_dict[specific_label][0] | |
| recyclable_advice = paper_item_num_dict[specific_label][1] | |
| elif material_label == 'plastic': | |
| specific_label, specific_label_idx, specific_probs = plastic_model.predict(inp) | |
| specific_preds = {name: prob for name, prob in zip(plastic_names, specific_probs.tolist())} | |
| specific_label = plastic_names[int(specific_label_idx)] | |
| recyclable_qn = plastic_item_num_dict[specific_label][0] | |
| recyclable_advice = plastic_item_num_dict[specific_label][1] | |
| elif material_label == 'glass': | |
| specific_label, specific_label_idx, specific_probs = glass_model.predict(inp) | |
| specific_preds = {name: prob for name, prob in zip(glass_names, specific_probs.tolist())} | |
| specific_label = glass_names[int(specific_label_idx)] | |
| recyclable_qn = glass_item_num_dict[specific_label][0] | |
| recyclable_advice = glass_item_num_dict[specific_label][1] | |
| elif material_label == 'metal': | |
| specific_label, specific_label_idx, specific_probs = metal_model.predict(inp) | |
| specific_preds = {name: prob for name, prob in zip(metal_names, specific_probs.tolist())} | |
| specific_label = metal_names[int(specific_label_idx)] | |
| recyclable_qn = metal_item_num_dict[specific_label][0] | |
| recyclable_advice = metal_item_num_dict[specific_label][1] | |
| elif material_label == 'others': | |
| specific_label, specific_label_idx, specific_probs = others_model.predict(inp) | |
| specific_preds = {name: prob for name, prob in zip(other_names, specific_probs.tolist())} | |
| specific_label = other_names[int(specific_label_idx)] | |
| recyclable_qn = others_item_num_dict[specific_label][0] | |
| recyclable_advice = others_item_num_dict[specific_label][1] | |
| return material_preds, specific_preds, recyclable_qn, recyclable_advice | |
| except: | |
| raise Exception("Invalid file format! Please only upload .jpg or .png files!") | |
| with gr.Blocks(title="Trash Classification", css="#custom_header {min-height: 3rem} #custom_title {min-height: 3rem; text-align: center}") as demo:#, theme=seafoam) as demo: | |
| gr.Markdown("# Recyclable Detector- Classification of trash and recyclables", elem_id="custom_title") | |
| gr.Markdown("Gradio Inference interface", elem_id="custom_title") | |
| with gr.Column(): | |
| with gr.Column(): | |
| with gr.Box(): | |
| gr.Markdown("## Inputs", elem_id="custom_header") | |
| input_image = gr.Image(label="Input Image") | |
| input_image.style(height=240) | |
| btn = gr.Button(value="Submit") | |
| btn.style(full_width=True) | |
| with gr.Column(): | |
| with gr.Box(): | |
| gr.Markdown("## Outputs", elem_id="custom_header") | |
| recycling_qn = gr.outputs.Textbox(label="Is this recyclable?") | |
| recycling_advice = gr.outputs.Textbox(label="It is not recyclable when:") | |
| with gr.Row(): | |
| material_probs = gr.outputs.Label(label="Material Prediction") | |
| item_probs = gr.outputs.Label(label="Item Prediction") | |
| gr.Examples( | |
| examples=examples, | |
| inputs=input_image, | |
| fn=predict_image, | |
| cache_examples=False, | |
| ) | |
| btn.click(predict_image, inputs=[input_image], | |
| outputs=[material_probs, item_probs, recycling_qn, recycling_advice]) | |
| if __name__ == "__main__": | |
| demo.launch(show_error=True) | |