Update app.py

app.py

@@ -1,9 +1,14 @@
 import gradio
+import numpy as np
 from transformers import ViTForImageClassification
 
 # Load the ViT model
 model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")
 
+# Define a function to convert a NumPy array to a torch tensor
+def numpy_to_tensor(array):
+  return torch.from_numpy(array).float()
+
 # Create a Gradio interface
 interface = gradio.Interface(
     fn=model,
@@ -13,5 +18,8 @@ interface = gradio.Interface(
     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()