Update app.py

app.py

@@ -10,15 +10,25 @@ class Prediction:
     prediction: TorchTensor
 
 app = fastapi.FastAPI(docs_url="/")
+from transformers import ViTForImageClassification
 
-# Load the pre-trained model
-pre_trained_model = torch.load('best_model.pth', map_location=torch.device('cpu'))
+# Define the number of classes in your custom dataset
+num_classes = 20
 
+# Initialize the ViTForImageClassification model
+model = ViTForImageClassification.from_pretrained(
+    'google/vit-base-patch16-224-in21k',
+    num_labels=num_classes  # Specify the number of classes
+)
+
+
+model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
 # Define a function to preprocess the input image
 def preprocess_input(input: fastapi.UploadFile):
     image = Image.open(input.file)
-    image = image.resize((224, 224))
+    image = image.resize((224, 224)).convert("RGB")
     input = np.array(image)
+    input = np.transpose(input, (2, 0, 1))
     input = torch.from_numpy(input).float()
     input = input.unsqueeze(0)
     return input
@@ -35,7 +45,17 @@ async def predict_endpoint(input:fastapi.UploadFile):
     prediction = model(input)
 
 
-    predicted_class = prediction.argmax(1).item()
-
-    # Return the predicted class in JSON format
-    return {"prediction": predicted_class}
+    logits = prediction.logits
+    num_top_predictions = 3
+    top_predictions = torch.topk(logits, k=num_top_predictions, dim=1)
+    
+    # Get the top 3 class indices and their probabilities
+    top_indices = top_predictions.indices.squeeze().tolist()
+    top_probabilities = torch.softmax(top_predictions.values, dim=1).squeeze().tolist()
+    
+    # Get the disease names for the top 3 predictions
+    disease_names = [disease_names[idx] for idx in top_indices]
+    
+    # Return the top 3 disease names and their probabilities in JSON format
+    response_data = [{"disease_name": name, "probability": prob} for name, prob in zip(disease_names, top_probabilities)]
+    return {"predictions": response_data}