Update app.py

app.py

@@ -1,24 +1,25 @@
 from transformers import pipeline
 import gradio as gr
 
-def select_model(model_name):
+def select_model(version):
+    if version == "v1":
+        model_name = "miittnnss/pet-classifier"
+    elif version == "v2":
+        model_name = "miittnnss/pet-classifier-v2"
     return pipeline("image-classification", model=model_name)
 
 def predict(image, model_name):
-    pipeline = select_model(model_name)
-    predicts = pipeline(image)
-    return image, {p["label"]: p["score"] for p in predicts}
+    pipeline_model = select_model(model_name)
+    predicts = pipeline_model(image)
+    return {p["label"]: p["score"] for p in predicts}
 
 iface = gr.Interface(
     predict,
     inputs=[
         gr.Image(label="Input", sources=["upload", "webcam"], type="pil"),
-        gr.Dropdown(label="Model Name", choices=["miittnnss/pet-classifier", "miittnnss/pet-classifier-v2"], value="miittnnss/pet-classifier-v2")
-    ],
-    outputs=[
-        gr.Image(label="Processed"),
-        gr.Label(label="Result")
+        gr.Radio(label="Model Version", choices=["v1", "v2"], value="v1")
     ],
+    outputs=gr.Label(label="Result"),
     title="Pet Classifier"
 )