app.py

import gradio as gr
from transformers import BeitFeatureExtractor, BeitForImageClassification
from PIL import Image
import requests
import numpy as np

# Load the pre-trained BEiT model and feature extractor
feature_extractor = BeitFeatureExtractor.from_pretrained('microsoft/beit-large-patch16-512')
model = BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-512')


def classify_image(input_image):
 image = Image.fromarray(input_image.astype('uint8'))
 inputs = feature_extractor(images=image, return_tensors="pt")
 outputs = model(**inputs)
 logits = outputs.logits
 predicted_class_idx = logits.argmax(-1).item()
 predicted_class = model.config.id2label[predicted_class_idx]
 return {"Predicted Class": predicted_class}


iface = gr.Interface(
 fn=classify_image,
 inputs=gr.inputs.Image(type="numpy"), # Specify input type as numpy array
 outputs="json",
 live=True,
 title="BEiT Classification",
 description="Upload an image and you will get a description"
)

if __name__ == "__main__":
 iface.launch()