Update app.py

app.py

@@ -1,5 +1,6 @@
-import gradio
+import gradio as gr
 import numpy as np
+import torch
 from transformers import ViTForImageClassification
 
 # Load the ViT model
@@ -7,19 +8,20 @@ model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224-i
 
 # Define a function to convert a NumPy array to a torch tensor
 def numpy_to_tensor(array):
-  return torch.from_numpy(array).float()
+    return torch.from_numpy(array).float()
+
+# Create input and output components
+input_component = gr.inputs.Image()
+output_component = gr.outputs.Label()
 
 # Create a Gradio interface
-interface = gradio.Interface(
+interface = gr.Interface(
     fn=model,
-    inputs="image",
-    outputs=["label"],
+    inputs=input_component,
+    outputs=output_component,
     title="ViT Image Classifier",
     description="This Gradio app allows you to classify images using a Vision Transformer (ViT) model."
 )
 
-# Set the input block to handle NumPy arrays
-interface.inputs[0].type = numpy_to_tensor
-
 # Launch the Gradio app
 interface.launch()