prediction function returns dictionary

app.py

@@ -74,15 +74,10 @@ def predict(img):
     with torch.no_grad():
         output = densenet(input_img)
 
-    predicted_probs = torch.sigmoid(output).to('cpu')
-    prediction = pd.DataFrame(predicted_probs, index=['predictions'],
-                          columns=labels).T.sort_values(by=['predictions'], ascending=False)
+    predicted_probs = torch.sigmoid(output).to('cpu').flatten().tolist()
+    prediction_dict = {labels[i]: predicted_probs[i] for i in range(len(labels)) if predicted_probs[i] > 0.5}
 
-
-    prediction_probs = prediction.query('predictions > 0.5').reset_index(names='morphs')['morphs'].to_list()
-    prediction_confidence = prediction.query('predictions > 0.5')['predictions'].to_list()
-
-    return prediction_probs, prediction_confidence
+    return prediction_dict