Update app.py

app.py

@@ -2,14 +2,17 @@ import gradio as gr
 import torch
 from transformers import ViTImageProcessor, ViTForImageClassification
 from PIL import Image
+import numpy as np
 
-# Load the pre-trained model and feature extractor
+# Load the pre-trained model and image processor
 model_name = "google/vit-base-patch16-224"
 image_processor = ViTImageProcessor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 
 # Define the prediction function
 def predict(image):
+    # Convert the dictionary to a PIL image
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
     inputs = image_processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
     logits = outputs.logits