Update app.py

app.py

@@ -6,6 +6,7 @@ from deepfake_vs_real import deepfake_classification
 from gym_workout_classification import workout_classification
 from augmented_waste_classifier import waste_classification
 from age_classification import age_classification
+from mnist_digits import classify_digit
 
 # Main classification function that calls the appropriate model based on selection.
 def classify(image, model_name):
@@ -23,37 +24,24 @@ def classify(image, model_name):
         return waste_classification(image)
     elif model_name == "age":
         return age_classification(image)
+    elif model_name == "mnist":
+        return classify_digit(image)
     else:
         return {"Error": "No model selected"}
 
 # Function to update the selected model and button styles.
 def select_model(model_name):
-    # Set each button's variant to "primary" if selected, otherwise "secondary"
-    gender_variant = "primary" if model_name == "gender" else "secondary"
-    emotion_variant = "primary" if model_name == "emotion" else "secondary"
-    dog_breed_variant = "primary" if model_name == "dog breed" else "secondary"
-    deepfake_variant = "primary" if model_name == "deepfake" else "secondary"
-    gym_workout_variant = "primary" if model_name == "gym workout" else "secondary"
-    waste_variant = "primary" if model_name == "waste" else "secondary"
-    age_variant = "primary" if model_name == "age" else "secondary"
-    # Return new state and update objects for each button in the specified order.
-    return (
-       model_name,
-       gr.update(variant=gender_variant),
-       gr.update(variant=emotion_variant),
-       gr.update(variant=dog_breed_variant),
-       gr.update(variant=deepfake_variant),
-       gr.update(variant=gym_workout_variant),
-       gr.update(variant=waste_variant),
-       gr.update(variant=age_variant)
-    )
+    model_variants = {
+        "gender": "secondary", "emotion": "secondary", "dog breed": "secondary", "deepfake": "secondary",
+        "gym workout": "secondary", "waste": "secondary", "age": "secondary", "mnist": "secondary"
+    }
+    model_variants[model_name] = "primary"
+    return (model_name, *(gr.update(variant=model_variants[key]) for key in model_variants))
 
 with gr.Blocks() as demo:
-    # Sidebar with title and model selection buttons.
     with gr.Sidebar():
         gr.Markdown("# SigLIP2 224")
         with gr.Row():
-            # Initialize buttons with variants. Default is "age" set to primary.
             age_btn = gr.Button("Age Classification", variant="primary")
             gender_btn = gr.Button("Gender Classification", variant="secondary")
             emotion_btn = gr.Button("Emotion Classification", variant="secondary")
@@ -61,44 +49,23 @@ with gr.Blocks() as demo:
             deepfake_btn = gr.Button("Deepfake vs Real", variant="secondary")
             gym_workout_btn = gr.Button("Gym Workout Classification", variant="secondary")
             waste_btn = gr.Button("Waste Classification", variant="secondary")
-
-        # State to hold the current model choice.
+            mnist_btn = gr.Button("MNIST Digit Classification", variant="secondary")
+        
         selected_model = gr.State("age")
-
         gr.Markdown("### Current Model:")
         model_display = gr.Textbox(value="age", interactive=False)
-        # Update display when state changes.
         selected_model.change(lambda m: m, selected_model, model_display)
 
-        # Set up click events for each button, updating state and button variants.
-        gender_btn.click(fn=lambda: select_model("gender"),
-                         inputs=[],
-                         outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-        emotion_btn.click(fn=lambda: select_model("emotion"),
-                          inputs=[],
-                          outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-        dog_breed_btn.click(fn=lambda: select_model("dog breed"),
-                            inputs=[],
-                            outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-        deepfake_btn.click(fn=lambda: select_model("deepfake"),
-                           inputs=[],
-                           outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-        gym_workout_btn.click(fn=lambda: select_model("gym workout"),
-                              inputs=[],
-                              outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-        waste_btn.click(fn=lambda: select_model("waste"),
-                        inputs=[],
-                        outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-        age_btn.click(fn=lambda: select_model("age"),
-                      inputs=[],
-                      outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
-
-    # Main interface: image input, analyze button, and prediction output.
+        buttons = [gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn, mnist_btn]
+        model_names = ["gender", "emotion", "dog breed", "deepfake", "gym workout", "waste", "age", "mnist"]
+        
+        for btn, name in zip(buttons, model_names):
+            btn.click(fn=lambda n=name: select_model(n), inputs=[], outputs=[selected_model] + buttons)
+    
     with gr.Column():
         image_input = gr.Image(type="numpy", label="Upload Image")
         analyze_btn = gr.Button("Classify / Predict")
         output_label = gr.Label(label="Prediction Scores")
-        # When the "Analyze" button is clicked, use the selected model to classify the image.
         analyze_btn.click(fn=classify, inputs=[image_input, selected_model], outputs=output_label)
 
 demo.launch()