Update app.py

app.py

@@ -3,7 +3,8 @@ from gender_classification import gender_classification
 from emotion_classification import emotion_classification
 from dog_breed import dog_breed_classification
 from deepfake_vs_real import deepfake_classification
-from gym_workout_classification import workout_classification
+from gym_workout_classification import gym_workout_classification
+from augmented_waste_classifier import waste_classification
 
 # Functions to update the model state when a button is clicked.
 def select_gender():
@@ -21,6 +22,9 @@ def select_deepfake():
 def select_gym_workout():
     return "gym workout"
 
+def select_waste():
+    return "waste"
+
 # Main classification function that calls the appropriate model based on selection.
 def classify(image, model_name):
     if model_name == "gender":
@@ -33,19 +37,22 @@ def classify(image, model_name):
         return deepfake_classification(image)
     elif model_name == "gym workout":
         return gym_workout_classification(image)
+    elif model_name == "waste":
+        return waste_classification(image)
     else:
         return {"Error": "No model selected"}
 
 with gr.Blocks() as demo:
     # Sidebar with title and model selection buttons.
     with gr.Sidebar():
-        gr.Markdown("# SigLIP2 Classification")
+        gr.Markdown("# SigLIP2 224")
         with gr.Row():
             gender_btn = gr.Button("Gender Classification")
             emotion_btn = gr.Button("Emotion Classification")
             dog_breed_btn = gr.Button("Dog Breed Classification")
             deepfake_btn = gr.Button("Deepfake vs Real")
             gym_workout_btn = gr.Button("Gym Workout Classification")
+            waste_btn = gr.Button("Waste Classification")
 
         # State to hold the current model choice.
         selected_model = gr.State("gender")
@@ -56,6 +63,7 @@ with gr.Blocks() as demo:
         dog_breed_btn.click(fn=select_dog_breed, inputs=[], outputs=selected_model)
         deepfake_btn.click(fn=select_deepfake, inputs=[], outputs=selected_model)
         gym_workout_btn.click(fn=select_gym_workout, inputs=[], outputs=selected_model)
+        waste_btn.click(fn=select_waste, inputs=[], outputs=selected_model)
 
         gr.Markdown("### Current Model:")
         model_display = gr.Textbox(value="gender", interactive=False)