Update app.py

app.py

@@ -67,9 +67,21 @@ async def predict_endpoint(input: fastapi.UploadFile):
 
     # Make a prediction
     prediction = model(input_data)
+    logits = prediction.logits
 
-
+    # Get the top N predictions
+    num_top_predictions = 20
+    top_predictions = torch.topk(logits, k=num_top_predictions, dim=1)
+    top_indices = top_predictions.indices.squeeze().tolist()
+    top_probabilities = torch.softmax(top_predictions.values, dim=1).squeeze().tolist()
     # Return the top N class names and their probabilities in JSON format
-    
+    top_diseases = [class_names[idx] for idx in top_indices]
+
+    # Create a response dictionary
+    response_data = {"predictions": []}
+
+    # Add the top N predictions and their probabilities to the response
+    for disease, probability in zip(top_diseases, top_probabilities):
+        response_data["predictions"].append(f"{disease}: {probability * 100:.2f}%")
 
-    return {"predictions": prediction}
+    return response_data