BCI/TrainingGenerator.py at master · JephDiel/BCI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import serial
import time

from ThinkGear import ThinkGear
import pygame
from pygame.locals import *
import random
from collections import deque
import numpy as np

my_device = ThinkGear("COM5")

pygame.init()

gameDisplay=pygame.display.set_mode((600,600))
pygame.display.set_caption("Platypus")

white=255,255,255
black=0,0,0
red=255,0,0

images = [[], [], [], []]


for i in range(0, 1):
    images[0] += [pygame.image.load("data/images/up/up" + str(i) + ".png")]
    images[1] += [pygame.image.load("data/images/down/down" + str(i) + ".png")]
    images[2] += [pygame.image.load("data/images/left/left" + str(i) + ".png")]
    images[3] += [pygame.image.load("data/images/right/right" + str(i) + ".png")]

reads_per_data = 256

raw_data = deque(maxlen=int(reads_per_data * 1)) #192
counter = 0
lastdata = 0

data_size = 180
time_per_data = 1

training_data_input = np.zeros((data_size * reads_per_data, reads_per_data))
training_data_output = np.zeros((data_size * reads_per_data, 4))

data_index = 0

for reads in range(reads_per_data):
        delay = time.time()
        my_device.fetch_data()
        data = my_device.data
        if 'eeg_raw' in data:
            raw_data.append(data['eeg_raw'])
            lastdata = data['eeg_raw']
        else:
            raw_data.append(lastdata)

for sample in range(data_size):
    dir = random.randrange(4)
    gameDisplay.blit(images[dir][random.randrange(len(images[dir]))], (0,0))
    pygame.display.update()
    time_save = time.time()
    for reads in range(reads_per_data):
        delay = time.time()
        my_device.fetch_data()
        data = my_device.data
        if 'eeg_raw' in data:
            raw_data.append(data['eeg_raw'])
            lastdata = data['eeg_raw']
        else:
            raw_data.append(lastdata)
        time.sleep(time_per_data/reads_per_data - (time.time() - delay))

    training_data_input[data_index] = list(raw_data)
    training_data_output[data_index][dir] = 1
    data_index += 1

    delay = time.time()

try:
    old_training_input = np.load("trainingdata/input.npy")
    old_training_output = np.load("trainingdata/output.npy")

    training_data_input = np.concatenate((training_data_input, old_training_input))
    training_data_output = np.concatenate((training_data_output, old_training_output))
except:
    print("Old Data Not Found, Overriding")
np.save("trainingdata/input.npy", training_data_input)
np.save("trainingdata/output.npy", training_data_output)

print("Added " + str(training_data_input.shape[0]) + " Items with " + str(training_data_input.shape[1]) + " elements each")
pygame.quit()
quit()