File size: 760 Bytes
1b2bb40
29e65a9
1b2bb40
1d612ff
 
 
 
 
1b2bb40
 
 
1d612ff
 
 
1b2bb40
 
 
 
 
1d612ff
1b2bb40
1d612ff
1b2bb40
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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)