Update main.py

main.py

@@ -19,35 +19,45 @@ 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.load_state_dict(torch.load("/kaggle/input/skincancer-vit/skin_cancer_model.pth", map_location=torch.device('cpu')))
 model.eval()
 
+# Define device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
 # Define class labels
 class_labels = ['benign_keratosis-like_lesions', 'basal_cell_carcinoma', 'actinic_keratoses', 'vascular_lesions', 'melanocytic_Nevi', 'melanoma', 'dermatofibroma']
 
+# Define optimal thresholds
+thresholds = [0.88134295, 0.43095806, 0.39622146, 0.90647435, 0.8128958, 0.05310565, 0.15926854]
+
 # Define image transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
 ])
 
-# Define API endpoint for model inference
+# Define API endpoint for model inference with class-specific thresholds
 @app.post('/predict')
 async def predict(file: UploadFile = File(...)):
     contents = await file.read()
     image = Image.open(io.BytesIO(contents))
-    image = transform(image).unsqueeze(0)  # Add batch dimension
+    image = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
     
     with torch.no_grad():
         outputs = model(image)
     
     # Calculate softmax probabilities
-    probabilities = torch.softmax(outputs.logits, dim=1)
+    probabilities = torch.softmax(outputs.logits, dim=1).cpu().numpy()[0]
     
     # Get predicted class index and its probability
-    predicted_idx = torch.argmax(probabilities).item()
+    predicted_idx = torch.argmax(torch.tensor(probabilities)).item()
     predicted_label = class_labels[predicted_idx]
-    predicted_accuracy = probabilities[0][predicted_idx].item()
+    predicted_probability = probabilities[predicted_idx]
     
-    return {'predicted_class': predicted_label, 'accuracy': predicted_accuracy}
+    # Check if the predicted probability meets the class-specific threshold
+    if predicted_probability < thresholds[predicted_idx]:
+        return {'predicted_class': 'uncertain', 'accuracy': predicted_probability}
+    else:
+        return {'predicted_class': predicted_label, 'accuracy': predicted_probability}