Update app.py

app.py

@@ -7,6 +7,8 @@ import gradio as gr
 from transformers import ViTFeatureExtractor
 from huggingface_hub import hf_hub_download
 import spaces
+from torchvision import transforms
+
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 model = None
@@ -17,6 +19,8 @@ VALID_DS_PATH = 'valid_ds.pth'
 valid_ds = torch.load(VALID_DS_PATH)
 
 
+
+
 from transformers import ViTModel
 from transformers.modeling_outputs import SequenceClassifierOutput
 import torch.nn as nn
@@ -68,6 +72,11 @@ def run_inference(image, device, valid_ds):
         model.eval()
         # feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k', do_rescale=False)
 
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),  # Resize to the model's input size
+        transforms.ToTensor(),
+    ])
+    
     image = Image.fromarray(image.astype('uint8'), 'RGB')
     input_tensor = transform(image)
     input_tensor = input_tensor.unsqueeze(0)  # Add a batch dimension