Pet-Classifier / app.py
miittnnss's picture
Update app.py
1d612ff verified
raw
history blame contribute delete
No virus
760 Bytes
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)