fix : make_prediction

image_ref/main.py

@@ -175,8 +175,8 @@ def make_prediction_duo(model, data, f_name, f_name2):
     for i in range(n_class):
         confidence_matrix[i]=np.mean(confidence_pred_list[i],axis=0)
 
-    df_cm = pd.DataFrame(cf_matrix / np.sum(cf_matrix, axis=1)[:, None], index=[i for i in range(2)],
-                         columns=['True','False'])
+    df_cm = pd.DataFrame(cf_matrix / np.sum(cf_matrix, axis=1)[:, None], index=[i for i in classes],
+                         columns=[i for i in classes])
     print('Saving Confusion Matrix')
     plt.clf()
     plt.figure(figsize=(14, 9))