new

app.py

@@ -1,33 +1,12 @@
-import requests
-
 import gradio as gr
-import torch
-from timm import create_model
-from timm.data import resolve_data_config
-from timm.data.tranforms import create_transform
-
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-
-
-extractor = AutoFeatureExtractor.from_pretrained("abhishek/autotrain_fashion_mnist_vit_base")
-model = AutoModelForImageClassification.from_pretrained("abhishek/autotrain_fashion_mnist_vit_base")
-transform = create_transform(
-    **resolve_data_config({}, model=model)
-)
-
-model.eval()
-
-def predict_fn(img):
-    img = img.convert('RGB')
-    img = transform(img).unsqueeze(0)
-
-    with torch.no_grad():
-        out = model(img)
-    
-    probabilites = torch.nn.functional.softmax(out[0], dim=0)
-
-    values, indices = torch.topk(probabilites, k=5)
-
-    return {LABELS[i]: v.item() for i, v in zip(indices, values)}
-
-gr.Interface(predict_fn, gr.inputs.Image(type='pil'), outputs='label').launch()
+from transformers import pipeline
+
+pipe = pipeline(task="image-classification", 
+                # model that can do 22k-category classification
+                model="abhishek/autotrain_fashion_mnist_vit_base")
+gr.Interface.from_pipeline(pipe, 
+                           title="22k Image Classification",
+                           description="Object Recognition using Microsoft BEIT",
+                           examples = ['wonder_cat.jpg', 'aki_dog.jpg',],
+                           article = "Author: <a href=\"https://huggingface.co/rowel\">Rowel Atienza</a>",
+                           ).launch()