Update app.py

app.py

@@ -1,6 +1,5 @@
 import torch
 import torchvision.transforms as transforms
-import torchvision.models as models
 import gradio as gr
 import numpy as np
 import tensorflow as tf
@@ -12,17 +11,17 @@ from transformers import ViTForImageClassification, ViTImageProcessor
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load ViT model from Hugging Face
-model_path = "andromeda01111/ViT_BCC"  # Ensure this file is uploaded to the Space
-state_dict = torch.load(model_path, map_location=device)
-vit_model.load_state_dict(state_dict)
-vit_model.to(device)
+# Load ViT model and processor from Hugging Face
+model_path = "andromeda01111/ViT_BCC"
+
+vit_model = ViTForImageClassification.from_pretrained(model_path).to(device)
 vit_model.eval()
+vit_processor = ViTImageProcessor.from_pretrained(model_path)
 
-# Load Neural Network model from Hugging Face
+# Load Neural Network model
 nn_model = tf.keras.models.load_model("andromeda01111/NN_BC")
 
-# Load scaler (ensure it's uploaded in the Hugging Face repo)
+# Load scaler
 scaler = joblib.load("scaler.pkl")
 
 # Class labels