Update app.py

app.py

@@ -70,18 +70,9 @@ async def predict_endpoint(input: fastapi.UploadFile):
     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]
+    predicted_class = torch.argmax(logits, dim=1).item()
 
     # Create a response dictionary
-    response_data = {"predictions": []}
+    return {"prediction": predicted_class}
 
-    # 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 response_data
+