HearingTestGUI.py

"""Architecture:
GUI for user
- Start test
- Play sounds
- Register button presses

Generate library of sounds
Play sounds
Record results

Score: hearing and playability

Display score

Save score and agregate
Dump data (for backup) after each test

Display agregated scores in second window
"""

import sounddevice as sd
import numpy as np
import time
from PyQt5 import QtCore, QtWidgets, QtGui, Qt
from PyQt5.QtWidgets import QMainWindow, QLabel, QGridLayout, QWidget, QTextEdit
from PyQt5.QtCore import QSize
import sys
from SoundLibrary import SoundLibrary
from HearingTest import HearingTest
import matplotlib
matplotlib.use('Qt5Agg')
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure


class HearingMplCanvas(FigureCanvas):
    """Ultimately, this is a QWidget (as well as a FigureCanvasAgg, etc.)."""

    false_color = np.array([0., 1., 0.])
    true_color = np.array([0., 0., 1.])

    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.baseline = -0.1
        fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = fig.add_subplot(111)

        # self.compute_initial_figure()

        FigureCanvas.__init__(self, fig)
        self.setParent(parent)
        fig.set_facecolor('k')
        FigureCanvas.setSizePolicy(self,
                                   QtWidgets.QSizePolicy.Expanding,
                                   QtWidgets.QSizePolicy.Expanding)
        FigureCanvas.updateGeometry(self)
        self.ul_bar = None
        self.ll_bar = None
        self.plot_result(HearingTest(parent.library, self))

    def plot_result(self, result: HearingTest):
        inds = np.arange(0, result.freqs.size)
        self.axes.set_facecolor('k')
        self.ul_bar = self.axes.bar(inds, result.upper_bounds - self.baseline,
                                    color=0.5*self.false_color, bottom=self.baseline)
        self.ll_bar = self.axes.bar(inds, result.lower_bounds - self.baseline,
                                    color=self.false_color, bottom=self.baseline)
        xlim = self.axes.get_xlim()
        for i in np.arange(result.lower_lim + self.baseline, result.upper_lim, result.block_size()):
            self.axes.plot(self.axes.get_xlim(), np.ones(2)*i, color='k')
        labels = [str(b) for b in result.freqs]
        labels.insert(0, '')
        self.axes.set_xticklabels(labels)
        self.axes.set_ylim(self.baseline, result.upper_lim - self.baseline)
        self.axes.set_xlim(xlim)

    def update_result(self, result: HearingTest) -> None:
        self.set_bar_heights(self.ul_bar, result.upper_bounds-self.baseline)
        self.set_bar_heights(self.ll_bar, result.lower_bounds-self.baseline)
        self.set_bar_colors([self.ul_bar, self.ll_bar], result.finished_freqs)
        self.draw()

    @staticmethod
    def set_bar_heights(bars: matplotlib.container.BarContainer, vals: np.ndarray) -> None:
        for bar, h in zip(bars, vals):
            bar.set_height(h)

    @staticmethod
    def set_bar_colors(bars, finished):
        depth = [0.5, 1]
        for i, bar_line in enumerate(bars):
            for bar, f in zip(bar_line, finished):
                if f:
                    bar.set_facecolor(depth[i] * HearingMplCanvas.true_color)
                else:
                    bar.set_facecolor(depth[i] * HearingMplCanvas.false_color)

    # TODO add animation


class HearingTestThread(QtCore.QThread):
    played_sound = QtCore.pyqtSignal(float, float, float)
    aborted = QtCore.pyqtSignal()

    def __init__(self, test: HearingTest):
        QtCore.QThread.__init__(self)
        self.test = test
        self.abort = False

    def __del__(self):
        self.wait()

    def run(self):
        finished = False
        while not finished:
            if self.abort:
                self.aborted.emit()
                break
            finished, freq, volume, played_time, next_sleep_time = self.test.play_next_sound()
            self.played_sound.emit(freq, volume, played_time)

            self.sleep(next_sleep_time)


class TestWindow(QMainWindow):

    start_text = "Press space to start..."

    def __init__(self, library):
        self.library = library
        self.test = None
        self.testing_thread = None
        self.test_running = False

        QMainWindow.__init__(self)

        self.setMinimumSize(QSize(640, 480))
        self.setWindowTitle("Hearing Test")
        self.setStyleSheet("QMainWindow { background-color: black }")

        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        grid_layout = QGridLayout(self)
        central_widget.setLayout(grid_layout)

        text_style = "QLabel { background-color: black; color: white; font: 24pt}"
        self.title = QLabel(start_text, self)
        self.title.setStyleSheet(text_style)
        self.title.setAlignment(QtCore.Qt.AlignCenter)

        self.log = QTextEdit()
        self.log.setReadOnly(True)
        self.log.setFocusPolicy(QtCore.Qt.NoFocus)
        self.log.hide()

        self.graph = HearingMplCanvas(parent=self)

        self.score_label = QLabel('', self)
        self.score_label.setStyleSheet(text_style)
        self.score_label.setAlignment(QtCore.Qt.AlignLeft)
        self.score_label.hide()

        grid_layout.addWidget(self.title, 0, 0)
        grid_layout.addWidget(self.log, 1, 0)
        grid_layout.addWidget(self.graph, 2, 0)
        grid_layout.addWidget(self.score_label, 3, 0)

    def keyPressEvent(self, event: QtGui.QKeyEvent):
        if event.key() == Qt.Qt.Key_D:
            if self.log.isHidden():
                self.log.show()
            else:
                self.log.hide()
            return
        if event.key() == Qt.Qt.Key_F:
            HearingTest.max_response_time = 0.01
            return
        if event.key() == Qt.Qt.Key_S:
            HearingTest.max_response_time = 1
            return
        if not self.test_running:
            if event.key() == Qt.Qt.Key_Escape:
                self.close()
            else:
                self.test_running = True
                self.run_test()
        else:
            if event.key() == Qt.Qt.Key_Escape:
                self.testing_thread.abort = True
            key_press_time = time.time()
            self.test.handle_key_press(key_press_time)
            self.record_key_press(event, key_press_time)

    def record_key_press(self, event: QtGui.QKeyEvent, _time):
        self.log.append('Key {} pressed at {}'.format(event.key(), _time))

    def run_test(self):
        self.log.setText('')
        self.title.setText('Test running')
        self.log.append('Starting...')
        self.score_label.hide()
        self.test = HearingTest(self.library, self.graph)
        self.testing_thread = HearingTestThread(self.test)
        self.testing_thread.played_sound.connect(self.sound_played)
        self.testing_thread.finished.connect(self.test_finished)
        self.testing_thread.aborted.connect(self.aborted)
        self.testing_thread.start()
        # need to prevent another thread starting

    def sound_played(self, freq, volume, played_time):
        self.log.append('Played freq {}, at volume {} and time {}'.format(freq, volume, played_time))

    def aborted(self):
        self.log.append('Test aborted')
        self.title.setText(self.start_text)
        self.test = None
        self.test_running = False

    def test_finished(self):
        if not self.test_running:
            return
        self.test_running = False
        self.title.setText(self.start_text)
        self.log.append("Finished test")
        self.log.append("Writing test result to file")
        thresholds = self.test.thresholds
        thresholds = thresholds.reshape(thresholds.shape[0], -1)  # convert to 2d
        with open('data.csv', 'ba') as file:
            np.savetxt(file, thresholds.T, fmt='%.2f', delimiter=',')
        # TODO calculate score
        # TODO display score
        score_str = self.test.get_score_text()
        self.score_label.setText(score_str)
        self.score_label.show()
        self.test = None


def main():
    device = sd.query_devices(sd.default.device['output'])
    fs = device['default_samplerate']
    length = 0.5
    f = [125, 250, 500, 750, 1000, 1500, 2000, 5000, 8000, 12000]

    library = SoundLibrary(fs, length, f)

    app = QtWidgets.QApplication(sys.argv)
    main_win = TestWindow(library)
    main_win.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()