Spaces:
Running
Running
import gradio as gr | |
from transformers import pipeline | |
def classify(image, model_name): | |
try: | |
pipe = pipeline("image-classification", model=model_name) | |
results = pipe(image) | |
return {result["label"]: round(result["score"], 2) for result in results} | |
except Exception as e: | |
return {"Error": str(e)} | |
# Gradio Blocks Interface | |
with gr.Blocks() as demo: | |
gr.Markdown( | |
""" | |
# Custom timm Model Image Classifier π | |
Explore the power of [timm](https://github.com/rwightman/pytorch-image-models) models for image classification using | |
the Hugging Face [Transformers pipeline](https://huggingface.co/docs/transformers/main_classes/pipelines). | |
With just a few lines of code, you can load any timm model hosted on the Hugging Face Hub and classify images effortlessly. | |
This application demonstrates how you can use the pipeline API to create a powerful yet minimalistic image classification tool. | |
## How to Use | |
1. Upload an image or use one of the provided examples. | |
2. Enter a valid timm model name from the Hugging Face Hub (e.g., `timm/resnet50.a1_in1k`). | |
3. View the top predictions and confidence scores! | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(): | |
image_input = gr.Image(type="pil", label="Upload an Image") | |
model_name_input = gr.Textbox( | |
label="Enter timm Model Name", | |
placeholder="e.g., timm/mobilenetv3_large_100.ra_in1k" | |
) | |
with gr.Column(): | |
output_label = gr.Label(num_top_classes=3, label="Top Predictions") | |
submit_button = gr.Button("Classify") | |
submit_button.click(fn=classify, inputs=[image_input, model_name_input], outputs=output_label) | |
gr.Examples( | |
examples=[ | |
["cat.jpg", "timm/mobilenetv3_small_100.lamb_in1k"], | |
["cat.jpg", "timm/resnet50.a1_in1k"], | |
], | |
inputs=[image_input, model_name_input] | |
) | |
gr.Markdown( | |
""" | |
## Learn More | |
- Check out the implementation in the `app.py` file to see how easy it is to integrate timm models. | |
- Dive into the [official blog post on timm integration](https://huggingface.co/blog/timm-transformers) for more insights. | |
""" | |
) | |
demo.launch() |