Update app.py

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import torch
+from transformers import AutoModel, AutoProcessor
 from gender_classification import gender_classification
 from emotion_classification import emotion_classification
 from dog_breed import dog_breed_classification
@@ -14,7 +16,7 @@ from alphabet_sign_language_detection import sign_language_classification
 from rice_leaf_disease import classify_leaf_disease
 from traffic_density import traffic_density_classification
 
-# Main classification function that calls the appropriate model based on selection.
+# Main classification function for multi-model classification.
 def classify(image, model_name):
     if model_name == "gender":
         return gender_classification(image)
@@ -58,42 +60,98 @@ def select_model(model_name):
     model_variants[model_name] = "primary"
     return (model_name, *(gr.update(variant=model_variants[key]) for key in model_variants))
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.Markdown("# Choose Domain")
-        with gr.Row():
-            age_btn = gr.Button("Age Classification", variant="primary")
-            gender_btn = gr.Button("Gender Classification", variant="secondary")
-            emotion_btn = gr.Button("Emotion Classification", variant="secondary")
-            dog_breed_btn = gr.Button("Dog Breed Classification", variant="secondary")
-            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")
-            mnist_btn = gr.Button("Digit Classify (0-9)", variant="secondary")
-            fashion_mnist_btn = gr.Button("Fashion MNIST Classification", variant="secondary")
-            food_btn = gr.Button("Indian/Western Food", variant="secondary")
-            bird_btn = gr.Button("Bird Species Classification", variant="secondary")
-            leaf_disease_btn = gr.Button("Rice Leaf Disease", variant="secondary")
-            sign_language_btn = gr.Button("Alphabet Sign Language", variant="secondary")
-            traffic_density_btn = gr.Button("Traffic Density", variant="secondary")
-        
-        selected_model = gr.State("age")
-        gr.Markdown("### Current Model:")
-        model_display = gr.Textbox(value="age", interactive=False)
-        selected_model.change(lambda m: m, selected_model, model_display)
+# Zero-Shot Classification Setup (SigLIP models)
 
-        buttons = [gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn, mnist_btn, fashion_mnist_btn, food_btn, bird_btn, leaf_disease_btn, sign_language_btn, traffic_density_btn]
-        model_names = ["gender", "emotion", "dog breed", "deepfake", "gym workout", "waste", "age", "mnist", "fashion_mnist", "food", "bird", "leaf disease", "sign language", "traffic density"]
-        
-        for btn, name in zip(buttons, model_names):
-            btn.click(fn=lambda n=name: select_model(n), inputs=[], outputs=[selected_model] + buttons)
+# Load the SigLIP models and processors
+sg1_ckpt = "google/siglip-so400m-patch14-384"
+siglip1_model = AutoModel.from_pretrained(sg1_ckpt, device_map="cpu").eval()
+siglip1_processor = AutoProcessor.from_pretrained(sg1_ckpt)
+
+sg2_ckpt = "google/siglip2-so400m-patch14-384"
+siglip2_model = AutoModel.from_pretrained(sg2_ckpt, device_map="cpu").eval()
+siglip2_processor = AutoProcessor.from_pretrained(sg2_ckpt)
+
+# Utilities for zero-shot classification.
+def postprocess_siglip(sg1_probs, sg2_probs, labels):
+    sg1_output = {labels[i]: sg1_probs[0][i].item() for i in range(len(labels))}
+    sg2_output = {labels[i]: sg2_probs[0][i].item() for i in range(len(labels))}
+    return sg1_output, sg2_output
+
+def siglip_detector(image, texts):
+    sg1_inputs = siglip1_processor(
+        text=texts, images=image, return_tensors="pt", padding="max_length", max_length=64
+    ).to("cpu")
+    sg2_inputs = siglip2_processor(
+        text=texts, images=image, return_tensors="pt", padding="max_length", max_length=64
+    ).to("cpu")
+    with torch.no_grad():
+        sg1_outputs = siglip1_model(**sg1_inputs)
+        sg2_outputs = siglip2_model(**sg2_inputs)
+        sg1_logits_per_image = sg1_outputs.logits_per_image
+        sg2_logits_per_image = sg2_outputs.logits_per_image
+        sg1_probs = torch.sigmoid(sg1_logits_per_image)
+        sg2_probs = torch.sigmoid(sg2_logits_per_image)
+    return sg1_probs, sg2_probs
+
+def infer(image, candidate_labels):
+    candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",")]
+    sg1_probs, sg2_probs = siglip_detector(image, candidate_labels)
+    return postprocess_siglip(sg1_probs, sg2_probs, labels=candidate_labels)
+
+# Build the Gradio Interface with two tab
+with gr.Blocks() as demo:
+    gr.Markdown("# Multi-Model & Zero-Shot Classification Interface")
     
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="numpy", label="Upload Image")
-            analyze_btn = gr.Button("Classify / Predict")
+    with gr.Tabs():
+        # Tab 1: Multi-Model Classification
+        with gr.Tab("Multi-Model Classification"):
+            with gr.Sidebar():
+                gr.Markdown("# Choose Domain")
+                with gr.Row():
+                    age_btn = gr.Button("Age Classification", variant="primary")
+                    gender_btn = gr.Button("Gender Classification", variant="secondary")
+                    emotion_btn = gr.Button("Emotion Classification", variant="secondary")
+                    dog_breed_btn = gr.Button("Dog Breed Classification", variant="secondary")
+                    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")
+                    mnist_btn = gr.Button("Digit Classify (0-9)", variant="secondary")
+                    fashion_mnist_btn = gr.Button("Fashion MNIST Classification", variant="secondary")
+                    food_btn = gr.Button("Indian/Western Food", variant="secondary")
+                    bird_btn = gr.Button("Bird Species Classification", variant="secondary")
+                    leaf_disease_btn = gr.Button("Rice Leaf Disease", variant="secondary")
+                    sign_language_btn = gr.Button("Alphabet Sign Language", variant="secondary")
+                    traffic_density_btn = gr.Button("Traffic Density", variant="secondary")
+                
+                selected_model = gr.State("age")
+                gr.Markdown("### Current Model:")
+                model_display = gr.Textbox(value="age", interactive=False)
+                selected_model.change(lambda m: m, selected_model, model_display)
+
+                buttons = [gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn, mnist_btn, fashion_mnist_btn, food_btn, bird_btn, leaf_disease_btn, sign_language_btn, traffic_density_btn]
+                model_names = ["gender", "emotion", "dog breed", "deepfake", "gym workout", "waste", "age", "mnist", "fashion_mnist", "food", "bird", "leaf disease", "sign language", "traffic density"]
+
+                for btn, name in zip(buttons, model_names):
+                    btn.click(fn=lambda n=name: select_model(n), inputs=[], outputs=[selected_model] + buttons)
+
+            with gr.Row():
+                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")
+                analyze_btn.click(fn=classify, inputs=[image_input, selected_model], outputs=output_label)
         
-        output_label = gr.Label(label="Prediction Scores")
-        analyze_btn.click(fn=classify, inputs=[image_input, selected_model], outputs=output_label)
+        # Tab 2: Zero-Shot Classification (SigLIP)
+        with gr.Tab("Zero-Shot Classification"):
+            gr.Markdown("## Compare SigLIP 1 and SigLIP 2 on Zero-Shot Classification")
+            with gr.Row():
+                with gr.Column():
+                    zs_image_input = gr.Image(type="pil", label="Upload Image")
+                    zs_text_input = gr.Textbox(label="Input a list of labels (comma separated)")
+                    zs_run_button = gr.Button("Run")
+                with gr.Column():
+                    siglip1_output = gr.Label(label="SigLIP 1 Output", num_top_classes=3)
+                    siglip2_output = gr.Label(label="SigLIP 2 Output", num_top_classes=3)
+            zs_run_button.click(fn=infer, inputs=[zs_image_input, zs_text_input], outputs=[siglip1_output, siglip2_output])
 
 demo.launch()