Update main.py

main.py

@@ -2,7 +2,7 @@ import os
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"  # Set cache directory to a writable location
 
 from fastapi import FastAPI, UploadFile, File
-from transformers import ViTForImageClassification, ViTFeatureExtractor, AutoTokenizer
+from transformers import ViTForImageClassification, ViTFeatureExtractor
 import torch
 import torch.nn as nn
 import torchvision.transforms as transforms
@@ -12,9 +12,16 @@ import io
 app = FastAPI()
 
 # Load the ViT model and its feature extractor
-model_name = "Anwarkh1/Skin_Cancer-Image_Classification"
+model_name = "google/vit-base-patch16-224-in21k"
 model = ViTForImageClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+
+# Load the trained model weights
+num_classes = 7
+model.classifier = nn.Linear(model.config.hidden_size, num_classes)
+# Load the trained weights
+model.load_state_dict(torch.load("skin_cancer_model.pth", map_location=torch.device('cpu')))
+model.eval()
 
 # Define class labels
 class_labels = ['benign_keratosis-like_lesions', 'basal_cell_carcinoma', 'actinic_keratoses', 'vascular_lesions', 'melanocytic_Nevi', 'melanoma', 'dermatofibroma']