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
+from transformers import ViTForImageClassification, ViTFeatureExtractor, AutoTokenizer
 import torch
 import torch.nn as nn
 import torchvision.transforms as transforms
@@ -12,16 +12,9 @@ import io
 app = FastAPI()
 
 # Load the ViT model and its feature extractor
-model_name = "google/vit-base-patch16-224-in21k"
+model_name = "Anwarkh1/Skin_Cancer-Image_Classification"
 model = ViTForImageClassification.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("models/Anwarkh1/Skin_Cancer-Image_Classification", map_location=torch.device('cpu')))
-model.eval()
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Define class labels
 class_labels = ['benign_keratosis-like_lesions', 'basal_cell_carcinoma', 'actinic_keratoses', 'vascular_lesions', 'melanocytic_Nevi', 'melanoma', 'dermatofibroma']