app.py

import gradio as gr
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification 
import requests
labels =  [
    "None",
    "Circle",
    "Triangle",
    "Square",
    "Pentagon",
    "Hexagon"
] 



feature_extractor = AutoImageProcessor.from_pretrained('0-ma/vit-geometric-shapes-tiny')
model = AutoModelForImageClassification.from_pretrained('0-ma/vit-geometric-shapes-tiny')

labels = []
def predict(image):
    feature_extractor = AutoImageProcessor.from_pretrained('0-ma/vit-geometric-shapes-tiny')
    model = AutoModelForImageClassification.from_pretrained('0-ma/vit-geometric-shapes-tiny')
    inputs = feature_extractor(images=[image], return_tensors="pt")
    logits = model(**inputs)['logits'].cpu().detach().numpy()
    predictions = np.argmax(logits, axis=1)    
    predicted_labels = [labels[prediction] for prediction in predictions]
    print(predicted_labels[0],logits[0][predictions[0]])


    return {"predicted_label" : predicted_labels[0] }
 
title = "Geometric Shape Classifier"
description = "A geometric shape setector."
examples = ['example/1_None.jpg','example/2_Circle.jpg','example/3_Triangle.jpg','example/4_Square.jpg','example/5_Pentagone.jpg','example/6_Hexagone.jpg']


gr.Interface(fn=predict,inputs=gr.Image(type="pil"),outputs=gr.Label(),title=title,description=description,examples=examples).launch()