Merge pull request #27 from enricobu96/feature/n-refactor

small refactor and some comments

main.py

@@ -5,8 +5,6 @@
 
 import argparse
 from enum import Enum
-from operator import sub
-
 from model.mediapipe.acquire_data_webcam import AcquireData
 from model.mediapipe.acquire_data_dataset import AcquireDataset
 from model.mediapipe.run_mediapipe import Run

model/kinect_depth/run_kinect.py

@@ -14,7 +14,6 @@
 class RunKinect:
 
     def run(self, classificationModel: str):
-
         """
         IMPORT MODEL
         Import model from binary dump
@@ -61,11 +60,17 @@ def run(self, classificationModel: str):
         g_ass = GestureAssistant(5, 60, 20, 0.8, True)
 
         while True:
-            # Get frame
+            """
+            GET FRAME
+            Get frame from IR camera
+            """
             ut_frame = user_tracker.read_frame()
             depth_frame = ut_frame.get_depth_frame()
 
-            # Image filtering
+            """
+            IMAGE FILTERING
+            Do image filtering on frames
+            """
             depth_frame_data = depth_frame.get_buffer_as_uint16()
             img = np.ndarray((depth_frame.height, depth_frame.width), dtype=np.uint16,
                          buffer=depth_frame_data).astype(np.float32)
@@ -76,6 +81,10 @@ def run(self, classificationModel: str):
                 img = (img - min_val) / (max_val - min_val)
             img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
 
+            """
+            SKELETON TRACKING
+            Start and do skeleton tracking (and get keypoints)
+            """
             if ut_frame.users:
                 for user in ut_frame.users:
                     if user.is_new():
@@ -87,7 +96,10 @@ def run(self, classificationModel: str):
                         row = pose_row
                         X = pd.DataFrame([row])
 
-                        # Do predictions
+                        """
+                        PREDICTIONS
+                        Do predictions
+                        """
                         if classificationModel == 'lr':
                             gesture_class, gesture_prob = self.__use_lr(model, X, img)
                         elif classificationModel == 'rc':
@@ -101,16 +113,18 @@ def run(self, classificationModel: str):
                         elif classificationModel == 'mlp':
                             gesture_class, gesture_prob = self.__use_cnn(model, X, img)
 
+                        """
+                        PRINT RESULTS ON SCREEN
+                        Just...this
+                        """
                         cv2.putText(img, 'CLASS', (95,12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
                         cv2.putText(img, gesture_class.split(' ')[0], (90,40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
                         cv2.putText(img, 'PROB', (15,12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
                         cv2.putText(img, str(round(gesture_prob[np.argmax(gesture_prob)],2)), (10,40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
 
                         if g_ass.addToBufferAndCheck(gesture_class, gesture_prob[np.argmax(gesture_prob)]):
                             print("sending..")
-                        socket.send(bytes(gesture_class,'utf-8')) #(byte?)
-
-                        print(gesture_class, gesture_prob)
+                        socket.send(bytes(gesture_class,'utf-8'))
 
 
             cv2.imshow("Depth", cv2.resize(img, (win_w, win_h)))
@@ -119,6 +133,10 @@ def run(self, classificationModel: str):
         k.close_camera()
         cv2.destroyAllWindows()
 
+    """
+    PREDICTION FUNCTIONS
+    One prediction function for every classification algorithm. Returns class and probability
+    """
     def __use_lr(self, model, X, img):
         gesture_class, gesture_prob = model.predict(X)[0], model.predict_proba(X)[0]
         return gesture_class, gesture_prob

model/mediapipe/run_mediapipe.py

@@ -72,7 +72,6 @@ def run(self, classificationModel: str):
                 results = holistic.process(image)
                 image.flags.writeable = True
 
-                # Draw pose landmarks
                 mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
                                       mp_drawing.DrawingSpec(
                                           color=(245, 117, 66), thickness=2, circle_radius=4),
@@ -81,13 +80,19 @@ def run(self, classificationModel: str):
                                       )
 
                 try:
-                    # Get pose keypoints
+                    """
+                    GET POSE KEYPOINTS
+                    Use mediapipe to get pose keypoints for prediction
+                    """
                     pose = results.pose_landmarks.landmark
                     pose_row = list(np.array([[landmark.x, landmark.y, landmark.z, landmark.visibility] for landmark in pose], ).flatten())    
                     row = pose_row
                     X = pd.DataFrame([row])
 
-                    # Do predictions
+                    """
+                    PREDICTIONS
+                    Do predictions
+                    """
                     if classificationModel == 'lr':
                         gesture_class, gesture_prob = self.__use_lr(model, X, image)
                     elif classificationModel == 'rc':
@@ -103,14 +108,16 @@ def run(self, classificationModel: str):
 
                     if g_ass.addToBufferAndCheck(gesture_class, gesture_prob[np.argmax(gesture_prob)]):
                         print("sending..")
-                        socket.send(bytes(gesture_class,'utf-8')) #(byte?)
+                        socket.send(bytes(gesture_class,'utf-8'))
 
+                    """
+                    PRINT RESULTS ON SCREEN
+                    Just...this
+                    """
                     cv2.putText(image, 'CLASS', (95,12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
                     cv2.putText(image, gesture_class.split(' ')[0], (90,40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
                     cv2.putText(image, 'PROB', (15,12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
                     cv2.putText(image, str(round(gesture_prob[np.argmax(gesture_prob)],2)), (10,40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
-
-                    print(gesture_class, gesture_prob)
 
                 except:
                     pass
@@ -121,6 +128,10 @@ def run(self, classificationModel: str):
         cv2.destroyAllWindows()
         k.close_camera()
 
+    """
+    PREDICTION FUNCTIONS
+    One prediction function for every classification algorithm. Returns class and probability
+    """
     def __use_lr(self, model, X, img):
         gesture_class, gesture_prob = model.predict(X)[0], model.predict_proba(X)[0]
         return gesture_class, gesture_prob

model/train.py

@@ -1,27 +1,48 @@
-from os import pipe
 import sys
 import pandas as pd
 from sklearn.metrics._plot.confusion_matrix import plot_confusion_matrix
 from sklearn.model_selection import train_test_split, learning_curve
 from sklearn.pipeline import make_pipeline
-from sklearn.preprocessing import StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScaler, Normalizer, QuantileTransformer, PowerTransformer
+from sklearn.preprocessing import StandardScaler
 from sklearn.linear_model import LogisticRegression, RidgeClassifier
-from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
+from sklearn.ensemble import RandomForestClassifier
 from sklearn.utils import shuffle
 from sklearn.svm import SVC
 from sklearn.neural_network import MLPClassifier
 import sklearn.metrics as skm
-from sklearn.metrics import precision_recall_fscore_support as score
 import pickle
 from matplotlib import pyplot as plt
 import numpy as np
 
+"""
+CLASSIFICATION ALGORITHM PARAMETERS
+Parameters for the different classification algorithms
+"""
+LR_C = 1e-2
+LR_MAX_ITER = 150
+RC_ALPHA = 1e2
+RF_N_ESTIMATORS = 30
+RF_MAX_DEPTH = 3
+RF_MIN_SAMPLES_LEAF = 10
+SVM_MAX_ITER = 200
+SVM_C = 1e-1
+MLP_ALPHA = 2e1
+MLP_HIDDEN_LAYER_SIZES = (32,)
+
+"""
+PLOT PARAMETERS
+Parameters for the learning curves plot
+"""
+LEARNING_CURVE_STEPS = 5
+CPU_THREADS = 22
+
 class Train:
 
     def train(self, is_nite: bool):
+
         """
         READ CSV AND PREPARE
-        Read the csv and prepare the model to predict
+        Read the csv and prepare the model(s) to predict
         """
         print('Loading dataset...', end='')
         sys.stdout.flush()
@@ -37,17 +58,17 @@ def train(self, is_nite: bool):
         sys.stdout.flush()
 
         """
-        TRAIN THE ML MODEL
-        Train the ML model, evaluate and serialize
+        TRAIN THE ML MODEL(S)
+        Train the ML model(s), evaluate and serialize
         """
         print('Training the model...')
         sys.stdout.flush()
         pipelines = {
-            'lr':make_pipeline(StandardScaler(), LogisticRegression(C=1e-2, max_iter=150)), # toccare C
-            'rc':make_pipeline(StandardScaler(), RidgeClassifier(alpha=1e2)), #toccare alpha
-            'rf':make_pipeline(StandardScaler(), RandomForestClassifier(n_estimators=30, max_depth=3, min_samples_leaf=10)),
-            'svm': make_pipeline(StandardScaler(), SVC(probability=True, max_iter=200, C=1e-1)),
-            'mlp': make_pipeline(StandardScaler(), MLPClassifier(alpha=2e1, random_state=42, hidden_layer_sizes=(32,)))
+            'lr':make_pipeline(StandardScaler(), LogisticRegression(C=LR_C, max_iter=LR_MAX_ITER)), 
+            'rc':make_pipeline(StandardScaler(), RidgeClassifier(alpha=RC_ALPHA)),
+            'rf':make_pipeline(StandardScaler(), RandomForestClassifier(n_estimators=RF_N_ESTIMATORS, max_depth=RF_MAX_DEPTH, min_samples_leaf=RF_MIN_SAMPLES_LEAF)),
+            'svm': make_pipeline(StandardScaler(), SVC(probability=True, max_iter=SVM_MAX_ITER, C=SVM_C)),
+            'mlp': make_pipeline(StandardScaler(), MLPClassifier(alpha=MLP_ALPHA, random_state=42, hidden_layer_sizes=MLP_HIDDEN_LAYER_SIZES))
         }
         fit_models = {}
         for alg, pipeline in pipelines.items():
@@ -58,29 +79,34 @@ def train(self, is_nite: bool):
             print('DONE')
             sys.stdout.flush()
 
-        print('Training the model...DONE')
+        print('Training the model(s)...DONE')
         sys.stdout.flush()
 
-        # Test the models
+        """
+        MODEL TESTING
+        Test the models and get stats and charts
+        """
+        print('Assessing model(s) accuracy...')
+        sys.stdout.flush()
+        # Accuracy
         self.test_accuracy(fit_models, X_test, y_test)
-
-
-        # Print confusion matrices
+        # Confusion matrices
         self.print_confusion_matrix('lr', fit_models['lr'], X_test, y_test)
         self.print_confusion_matrix('rc', fit_models['rc'], X_test, y_test)
         self.print_confusion_matrix('rf', fit_models['rf'], X_test, y_test)
         self.print_confusion_matrix('svm', fit_models['svm'], X_test, y_test)
         self.print_confusion_matrix('mlp', fit_models['mlp'], X_test, y_test)
-
-        # Print learning curves
+        # Learning curves
         self.print_learning_curves('lr', fit_models['lr'], X_test, y_test, X_train, y_train)
         self.print_learning_curves('rc', fit_models['rc'], X_test, y_test, X_train, y_train)
         self.print_learning_curves('rf', fit_models['rf'], X_test, y_test, X_train, y_train)
         self.print_learning_curves('svm', fit_models['svm'], X_test, y_test, X_train, y_train)
         self.print_learning_curves('mlp', fit_models['mlp'], X_test, y_test, X_train, y_train)
 
-
-        # Serialize models
+        """
+        SERIALIZE MODELS
+        Save trained model(s) into pickle files
+        """
         if not is_nite:
             for k in pipelines.keys():
                 fileName = ('./model/mediapipe/prediction_models/prediction_model_' + k + '.pkl')
@@ -92,6 +118,10 @@ def train(self, is_nite: bool):
                 with open(fileName, 'wb') as f:
                     pickle.dump(fit_models[k], f)
 
+    """
+    TEST ACCURACY
+    Get accuracy, precision, recall, F1 score and support for every class in every model. Save results into file
+    """
     def test_accuracy(self, fit_models, X_test, y_test):
         out = ''
         for alg, model in fit_models.items():
@@ -102,19 +132,27 @@ def test_accuracy(self, fit_models, X_test, y_test):
         with open('./extra/test_reports/test_reports.txt', mode='w', newline='') as f:
             f.write(out)
 
+    """
+    CONFUSION MATRIX
+    Calculate confusion matrices for every class in every model. Save results into file(s)
+    """
     def print_confusion_matrix(self, alg, model, X_test, y_test):
         fig = plt.figure()
-        # xy = ['balc', 'bend', 'boxx', 'clap', 'marc', 'onew', 'wave'] # per APE
-        xy = ['dab', 'tp', 'rarmm', 'rarmt', 'larmm', 'larmt', 'st'] # per Custom
+        # xy = ['balc', 'bend', 'boxx', 'clap', 'marc', 'onew', 'wave'] # x,y labels for APE dataset
+        xy = ['dab', 'tp', 'rarmm', 'rarmt', 'larmm', 'larmt', 'st'] # x,y labels for custom dataset
         y_predict = model.predict(X_test)
         c_matrix = skm.confusion_matrix(y_test, y_predict)
         disp = plot_confusion_matrix(model, X_test, y_test, display_labels=xy, cmap='plasma')
         plt.xlabel('Predicted label')
         plt.savefig('./extra/test_reports/conf_matrix_' + alg + '.png')
 
+    """
+    LEARNING CURVES
+    Compute learning curves for every model. Save results into file(s)
+    """
     def print_learning_curves(self, alg, model, X_test, y_test, X_train, y_train):
         fig = plt.figure()
-        train_sizes, train_scores, validation_scores = learning_curve(model, X_train, y_train, train_sizes=np.linspace(0.1,1.0,5), n_jobs=22)
+        train_sizes, train_scores, validation_scores = learning_curve(model, X_train, y_train, train_sizes=np.linspace(0.1,1.0,LEARNING_CURVE_STEPS), n_jobs=CPU_THREADS)
         train_scores_mean = train_scores.mean(axis = 1)
         validation_scores_mean = validation_scores.mean(axis = 1)
         plt.plot(train_sizes, train_scores_mean, train_sizes, validation_scores_mean)

test/client.py

@@ -2,7 +2,6 @@
 import time
 import zmq
 
-
 signal.signal(signal.SIGINT, signal.SIG_DFL)
 
 context = zmq.Context()