Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import pipeline | |
import tensorflow | |
import torch | |
global pipeline | |
global default_model_name | |
default_model_name = "google/vit-base-patch16-224" | |
def predict(image, model_name): | |
if model_name == "": | |
model_name = default_model_name | |
pipe = pipeline(task="image-classification", model=model_name) | |
predictions = pipe(image) | |
return {p["label"]: p["score"] for p in predictions} | |
with gr.Blocks() as demo: | |
with gr.Row(): | |
gr.Markdown( | |
""" | |
# Settings | |
[Here](https://huggingface.co/models?pipeline_tag=image-classification&sort=downloads) are some popular image classification models. | |
Or use default model **"google/vit-base-patch16-224"** | |
""") | |
gr.Markdown( | |
""" | |
# Image Classifier Result | |
""") | |
with gr.Row(): | |
with gr.Column(scale=1): | |
#input_model = gr.Textbox(label="Enter a custom model name:", value=default_model_name, scale=1) | |
input_model = gr.Textbox(label="Enter a custom model name:", scale=1) | |
gr.Markdown("Upload image") | |
#images_input = gr.File(file_count="multiple", file_types=["image"], label="Input images", scale=1) | |
#images_input = gr.Files(file_count="multiple", file_types=["image"], label="Input images", scale=1) | |
input_image = gr.Image(label="Input Image", type="filepath") | |
#output = gr.Label(label="Output", num_top_classes=3, scale=2) | |
output = gr.Label(num_top_classes=10, scale=2) | |
with gr.Row(equal_height=True): | |
clear_button = gr.ClearButton(value="Clear", scale=0) | |
submit_button = gr.Button(value="Submit", variant="primary", scale=0) | |
submit_button.click(fn=predict, inputs=[input_image, input_model], outputs=output) | |
clear_button.click(lambda: [None, None, None], outputs=[input_model, input_image, output]) | |
demo.launch() | |