Update app.py

app.py

@@ -1,24 +1,27 @@
 import gradio as gr
-import numpy as np
 import torch
-from transformers import ViTForImageClassification
+from transformers import ViTFeatureExtractor, ViTForImageClassification
 
 # Load the ViT model
-model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
+model_name = "google/vit-base-patch16-224-in21k"
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
 
-# Define a function to convert a NumPy array to a torch tensor
-def numpy_to_tensor(array):
-    return torch.from_numpy(array).float()
-
-# Create input and output components
-input_component = gr.inputs.Image()
-output_component = gr.outputs.Label()
+# Define a function to preprocess the image and classify it
+def classify_image(input_image):
+    # Preprocess the image using the feature extractor
+    inputs = feature_extractor(input_image, return_tensors="pt")
+    # Perform inference with the model
+    outputs = model(**inputs)
+    # Get the predicted label
+    predicted_class = torch.argmax(outputs.logits, dim=1).item()
+    return predicted_class
 
 # Create a Gradio interface
 interface = gr.Interface(
-    fn=model,
-    inputs=input_component,
-    outputs=output_component,
+    fn=classify_image,
+    inputs=gr.inputs.Image(type="numpy", label="Upload an image"),
+    outputs="label",
     title="ViT Image Classifier",
     description="This Gradio app allows you to classify images using a Vision Transformer (ViT) model."
 )