From 5529d7eef3230b1cdd61acce06ce6d8cb916d8a8 Mon Sep 17 00:00:00 2001
From: Hevia <anthony@hevia.dev>
Date: Sat, 21 Nov 2020 22:50:11 -0500
Subject: [PATCH] cleaned up python code

---
 hardware/common.py           | 13 ++++++++++++-
 hardware/hardware_starter.py | 25 +++++++++++--------------
 hardware/model_helper.py     |  4 ++--
 hardware/train.py            |  2 +-
 4 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/hardware/common.py b/hardware/common.py
index b974015..2b42a26 100644
--- a/hardware/common.py
+++ b/hardware/common.py
@@ -1,3 +1,14 @@
+import json
+
 ## Used for printing statements to application console
 def debugPrint(message):
-  print('BACKGROUND DEBUG PRINT:', message)
\ No newline at end of file
+  print('BACKGROUND DEBUG PRINT:', message)
+
+## Call this function to send the predicted emotion to the desktop
+## Calling this function will end this process
+def sendPredictedEmotion(emotion: str) -> None:
+  print(json.dumps({ 'emotion' : emotion }))
+
+## Call to signal to desktop that a connection has been established
+def confirmConnection() -> None:
+  print(json.dumps({ 'hasConfirmed' : True }))
\ No newline at end of file
diff --git a/hardware/hardware_starter.py b/hardware/hardware_starter.py
index 87d34c1..d71fd55 100644
--- a/hardware/hardware_starter.py
+++ b/hardware/hardware_starter.py
@@ -1,25 +1,22 @@
 import time
 import json
 import random
-from common import debugPrint
+from common import debugPrint, sendPredictedEmotion, confirmConnection
 from LSL_helper import connectToEEG, recordEEG
 from model_helper import predict_emotion
 
-
-## Call this function to send the predicted emotion to the desktop
-## Calling this function will end this process
-def sendPredictedEmotion(emotion: str) -> None:
-  print(json.dumps({ 'emotion' : emotion }))
-
-## Call to signal to desktop that a connection has been established
-def confirmConnection() -> None:
-  print(json.dumps({ 'hasConfirmed' : True }))
-
-# TODO: Error handling for None
-# TODO: Add remaining model code
+# The main EEG script that the desktop calls
 stream_connection = connectToEEG()
+if stream_connection == None:
+	debugPrint("Could not connect to EEG stream... disconnecting...")
+	exit()
+
 confirmConnection()
 data = recordEEG(stream_connection)
+if data == None:
+	debugPrint("Could not record data from EEG headset... exiting.....")
+	exit()
+
 emotion = predict_emotion(data)
-debugPrint(f"the emotion was: {emotion}")
+debugPrint(f"The predicted emotion was: {emotion}")
 sendPredictedEmotion(emotion)
\ No newline at end of file
diff --git a/hardware/model_helper.py b/hardware/model_helper.py
index b9ae2db..36a9286 100644
--- a/hardware/model_helper.py
+++ b/hardware/model_helper.py
@@ -11,9 +11,9 @@
 except ImportError:
 	debugPrint("Unable to import Tensorflow! Is it installed?")
 
-def predict_emotion(data: deque) -> None:
+def predict_emotion(data: deque) -> str:
 	debugPrint("Loading model....")
-	model_path = f"{Path.cwd()}/model.h5"
+	model_path = f"{Path.cwd()}/checkpoints/model.h5"
 	cpu = False
 	model_path
 	if cpu:
diff --git a/hardware/train.py b/hardware/train.py
index 2f2deea..133866a 100644
--- a/hardware/train.py
+++ b/hardware/train.py
@@ -40,4 +40,4 @@ def build_model():
 
 def train_model(model, train_X, train_Y, batch_size, epochs):
 	model.fit(train_X, train_y, batch_size=batch_size, epochs=epochs)
-	model.save("./model.h5")
+	model.save("./checkpoints/model.h5")