app.py

# 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)