diff --git a/image_ref/main.py b/image_ref/main.py
index 12a696e6222cc3fb470b25081be23b9678df130c..767cce8e7ce199cc67127807ee3940c006154f86 100644
--- a/image_ref/main.py
+++ b/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))