Diff of /python/render.py [000000] .. [a32498]

Switch to unified view

a b/python/render.py
1
#!/usr/bin/python
2
# Renders a window with graph values for each sensor and a box for gyro values.
3
try:
4
    import psyco
5
    psyco.full()
6
except ImportError:
7
    print('No psyco. Expect poor performance. Not really...')
8
import platform
9
import sys
10
import time
11
12
import pygame
13
from pygame import FULLSCREEN
14
15
from emokit.emotiv import Emotiv
16
from emokit.packet import EmotivExtraPacket
17
from emokit.util import get_quality_scale_level_color
18
19
if platform.system() == "Windows":
20
    pass
21
22
23
class Grapher(object):
24
    """
25
    Worker that draws a line for the sensor value.
26
    """
27
28
    def __init__(self, screen, name, i, old_model=False):
29
        """
30
        Initializes graph worker
31
        """
32
        self.screen = screen
33
        self.name = name
34
        self.range = float(1 << 13)
35
        self.x_offset = 40
36
        self.y = i * gheight
37
        self.buffer = []
38
        font = pygame.font.Font(None, 24)
39
        self.text = font.render(self.name, 1, (255, 0, 0))
40
        self.text_pos = self.text.get_rect()
41
        self.text_pos.centery = self.y + gheight
42
        self.first_packet = True
43
        self.y_offset = 0
44
        self.old_model = old_model
45
46
    def update(self, packet):
47
        """
48
        Appends value and quality values to drawing buffer.
49
        """
50
        if len(self.buffer) == 800 - self.x_offset:
51
            self.buffer = self.buffer[1:]
52
        self.buffer.append([packet.sensors[self.name]['value'], packet.sensors[self.name]['quality']])
53
54
    def calc_y(self, val):
55
        """
56
        Calculates line height from value.
57
        """
58
        return val - self.y_offset + gheight
59
60
    def draw(self):
61
        """
62
        Draws a line from values stored in buffer.
63
        """
64
        if len(self.buffer) == 0:
65
            return
66
67
        if self.first_packet:
68
            self.y_offset = self.buffer[0][0]
69
            # print(self.y_offset)
70
            self.first_packet = False
71
        pos = self.x_offset, self.calc_y(self.buffer[0][0]) + self.y
72
        for i, (value, quality) in enumerate(self.buffer):
73
            y = self.calc_y(value) + self.y
74
            color = get_quality_scale_level_color(quality, self.old_model)
75
            pygame.draw.line(self.screen, color, pos, (self.x_offset + i, y))
76
            pos = (self.x_offset + i, y)
77
        self.screen.blit(self.text, self.text_pos)
78
79
80
def main():
81
    """
82
    Creates pygame window and graph drawing workers for each sensor.
83
    """
84
    global gheight
85
    pygame.init()
86
    screen = pygame.display.set_mode((800, 600))
87
    graphers = []
88
    recordings = []
89
    recording = False
90
    record_packets = []
91
    updated = False
92
    cursor_x, cursor_y = 400, 300
93
    fullscreen = False
94
    with Emotiv(display_output=False) as emotiv:
95
        for name in 'AF3 F7 F3 FC5 T7 P7 O1 O2 P8 T8 FC6 F4 F8 AF4'.split(' '):
96
            graphers.append(Grapher(screen, name, len(graphers), emotiv.old_model))
97
        while emotiv.running:
98
            for event in pygame.event.get():
99
                if event.type == pygame.QUIT:
100
                    emotiv.close()
101
                    return
102
                if event.type == pygame.KEYDOWN:
103
                    if event.key == pygame.K_ESCAPE:
104
                        emotiv.close()
105
                        return
106
                    elif event.key == pygame.K_f:
107
                        if fullscreen:
108
                            screen = pygame.display.set_mode((800, 600))
109
                            fullscreen = False
110
                        else:
111
                            screen = pygame.display.set_mode((800, 600), FULLSCREEN, 16)
112
                            fullscreen = True
113
                    elif event.key == pygame.K_r:
114
                        if not recording:
115
                            record_packets = []
116
                            recording = True
117
                        else:
118
                            recording = False
119
                            recordings.append(list(record_packets))
120
                            record_packets = None
121
            packets_in_queue = 0
122
            try:
123
                while packets_in_queue < 8:
124
                    packet = emotiv.dequeue()
125
126
                    if packet is not None:
127
                        if type(packet) != EmotivExtraPacket:
128
                            if abs(packet.sensors['X']['value']) > 1:
129
                                cursor_x = max(0, min(cursor_x, 800))
130
                                cursor_x -= packet.sensors['X']['value']
131
                            if abs(packet.sensors['Y']['value']) > 1:
132
                                cursor_y += packet.sensors['Y']['value']
133
                                cursor_y = max(0, min(cursor_y, 600))
134
                            map(lambda x: x.update(packet), graphers)
135
                            if recording:
136
                                record_packets.append(packet)
137
                            updated = True
138
                            packets_in_queue += 1
139
                    time.sleep(0.001)
140
            except Exception as ex:
141
                print("EmotivRender DequeuePlotError ", sys.exc_info()[0], sys.exc_info()[1], sys.exc_info()[2],
142
                      " : ", ex)
143
144
            if updated:
145
                screen.fill((75, 75, 75))
146
                map(lambda x: x.draw(), graphers)
147
                pygame.draw.rect(screen, (255, 255, 255), (cursor_x - 5, cursor_y - 5, 10, 10), 0)
148
                pygame.display.flip()
149
                updated = False
150
151
152
if __name__ == "__main__":
153
    try:
154
        gheight = 580 // 14
155
        main()
156
    except Exception as ex:
157
        print("EmotivRender ", sys.exc_info()[0], sys.exc_info()[1], sys.exc_info()[2], " : ", ex)