Simplify code by removing repetition

1 files changed, 23 insertions, 46 deletions

diff --git a/train.py b/train.py
index 349d494..ef2df18 100755
--- a/train.py
+++ b/train.py
@@ -63,18 +63,17 @@ def test_split(n_faces, raw_faces, split, seed):
     faces_test = faces_test.reshape(n_faces*n_test_faces, n_pixels)
     return faces_train, faces_test, target_train, target_test
 
-def draw_conf_mat(args, target_test, target_pred):
+def draw_results(args, target_test, target_pred):
     acc_sc = accuracy_score(target_test, target_pred)
-    if not args.classifyalt:
-        cm = confusion_matrix(target_test, target_pred)
-        print('Accuracy: ', acc_sc)
-        if (args.conf_mat):
-            plt.matshow(cm, cmap='Blues')
-            plt.colorbar()
-            plt.ylabel('Actual')
-            plt.xlabel('Predicted')
-            plt.show()
-    return acc_sc
+    cm = confusion_matrix(target_test, target_pred)
+    print('Accuracy: ', acc_sc)
+    if (args.conf_mat):
+        plt.matshow(cm, cmap='Blues')
+        plt.colorbar()
+        plt.ylabel('Actual')
+        plt.xlabel('Predicted')
+        plt.show()
+    return
 
 def test_model(M, faces_train, faces_test, target_train, target_test, args):
     raw_faces_train = faces_train
@@ -185,8 +184,7 @@ def test_model(M, faces_train, faces_test, target_train, target_test, args):
         plt.show()
 
         #Better n_neighbors = 2
-    accuracy = draw_conf_mat(args, target_test, target_pred)
-    return target_pred, accuracy, distances
+    return target_pred, distances
 
 def main():
     parser = argparse.ArgumentParser()
@@ -233,32 +231,23 @@ def main():
         faces_train = faces_train.reshape(n_faces, int(faces_train.shape[0]/n_faces), n_pixels)
         target_train = target_train.reshape(n_faces, int(target_train.shape[0]/n_faces))
 
-        accuracy = np.zeros(n_faces)
         distances = np.zeros((n_faces, faces_test.shape[0]))
         for i in range(n_faces):
-            target_pred, accuracy[i], distances[i] = test_model(args.eigen, faces_train[i], faces_test, target_train[i], target_test, args)
+            target_pred, distances[i] = test_model(args.eigen, faces_train[i], faces_test, target_train[i], target_test, args)
         target_pred = np.argmin(distances, axis=0)
-        acc_sc = accuracy_score(target_test, target_pred)
-        cm = confusion_matrix(target_test, target_pred)
-        print('Total Accuracy: ', acc_sc)
-        if (args.conf_mat):
-            plt.matshow(cm, cmap='Blues')
-            plt.colorbar()
-            plt.ylabel('Actual')
-            plt.xlabel('Predicted')
-            plt.show()
-        return
+    elif args.reigen:
+        target_pred = np.zeros((args.reigen-args.eigen, 2*n_faces))
+        accuracy    = np.zeros((args.reigen-args.eigen, 2*n_faces))
+        rec_error   = np.zeros((args.reigen-args.eigen, 2*n_faces))
 
-    if args.reigen:
-        accuracy = np.zeros(args.reigen-args.eigen)
-        rec_error = np.zeros((args.reigen-args.eigen, 2*n_faces))
         for M in range(args.eigen, args.reigen):
             start = timer()
-            target_pred, accuracy[M - args.eigen], rec_error[M - args.eigen] = test_model(M, faces_train, faces_test, target_train, target_test, args)
+            target_pred[i], rec_error[M - args.eigen] = test_model(M, faces_train, faces_test, target_train, target_test, args)
             end = timer()
             print("Run with", M, "eigenvalues completed in ", end-start, "seconds")
             print("Memory Used:", psutil.Process(os.getpid()).memory_info().rss)
-        #plot
+            accuracy[i] = accuracy_score(target_test, target_pred[i])
+        # Plot
         print('Max efficiency of ', max(accuracy), '% for M =', np.argmax(accuracy))
         plt.plot(range(args.eigen, args.reigen), 100*accuracy)
         plt.xlabel('Number of Eigenvectors used (M)')
@@ -266,34 +255,22 @@ def main():
         plt.grid(True)
         plt.show()
     elif args.ensemble:
-        accuracy = np.zeros(args.ensemble)
         rec_error = np.zeros((args.ensemble, n_faces, faces_test.shape[0]))
         target_pred = np.zeros((args.ensemble, target_test.shape[0]))
         for i in range(args.ensemble):
-            target_pred[i], accuracy[i], rec_error[i] = test_model(args.eigen, faces_train_bagged[i], faces_test, target_train, target_test, args)
+            target_pred[i], rec_error[i] = test_model(args.eigen, faces_train_bagged[i], faces_test, target_train, target_test, args)
 
         target_pred_comb = np.zeros(target_pred.shape[1])
         target_pred = target_pred.astype(int).T
         for i in range(target_pred.shape[0]):
             target_pred_comb[i] = np.bincount(target_pred[i]).argmax()
-        print(target_pred_comb)
-        acc_sc = accuracy_score(target_test, target_pred_comb)
-        cm = confusion_matrix(target_test, target_pred_comb)
-        print('Total Accuracy: ', acc_sc)
-        if (args.conf_mat):
-            plt.matshow(cm, cmap='Blues')
-            plt.colorbar()
-            plt.ylabel('Actual')
-            plt.xlabel('Predicted')
-            plt.show()
-        return
-
-
+        target_pred = target_pred_comb
     else:
         M = args.eigen
         start = timer()
-        test_model(M, faces_train, faces_test, target_train, target_test, args)
+        target_pred, distances = test_model(M, faces_train, faces_test, target_train, target_test, args)
         end = timer()
+    draw_results(args, target_test, target_pred)
 
 if __name__ == "__main__":
     main()