from transformers import pipeline
import gradio as gr

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_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.Radio(label="Model Version", choices=["v1", "v2"], value="v1")
    ],
    outputs=gr.Label(label="Result"),
    title="Pet Classifier"
)

iface.launch(debug=True)