conan/uiwidget/widgetspeaking.py

import numpy as np
import pyqtgraph as pg
from PyQt5 import QtWidgets, QtCore
from utils.util import get_circle


class WidgetSpeaking(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(WidgetSpeaking, self).__init__(parent)

        layout = QtWidgets.QGridLayout()

        self.measures = QtWidgets.QWidget()
        self.measuresLayout = QtWidgets.QHBoxLayout()

        # Setup Speaking Graph Plot Widget
        self.speakingGraph = pg.PlotWidget()
        self.speakingGraph.setBackground('w')
        self.speakingGraph.hideAxis('left')
        self.speakingGraph.hideAxis('bottom')
        self.speakingGraph.setAspectLocked()
        self.speakingGraph.setYRange(-2.25, 2.25, padding=0)
        self.speakingGraph.setXRange(-2.25, 2.25, padding=0)
        self.speakingGraph.getPlotItem().setTitle(title='Speaking Distribution')

        self.speakingGraphPlots = []

        layout.addWidget(self.speakingGraph, 1, 0, 1, 1)
        layout.addWidget(self.measures, 1, 1, 1, 2)
        layout.setColumnStretch(0, 1)
        layout.setColumnStretch(1, 2)

        self.setLayout(layout)
        self.colors = None
        self.labels = dict()
        self.speakBlob = None
        self.positions = dict()
        self.idBlobs = dict()

    @QtCore.pyqtSlot(dict, int)
    def updateLables(self, speak, numberIDs):
        """ Update labels when segment was changed"""
        for id_no in range(numberIDs):
            self.labels['RST'][id_no].setText('%.2f' % speak[id_no][1])
            # self.labels['Turns'][id_no].setText('%i' % speak[id_no][2])
            # self.labels['TurnLength'][id_no].setText('%.2f sec' % speak[id_no][3])
            # self.labels['TurnPerMin'][id_no].setText('%.2f /min' % speak[id_no][4])

    @QtCore.pyqtSlot(dict, list, int)
    def setInit(self, speak, colors, numberIDs):
        self.colors = colors
        idLayout = QtWidgets.QVBoxLayout()

        labelID = QtWidgets.QLabel('      ')
        labelID.setFixedWidth(60)
        labelID.setFixedHeight(20)
        labelRST = QtWidgets.QLabel('Relative speaking time:')
        labelTurn = QtWidgets.QLabel('Number of speaking turns:')
        labelTurnLen = QtWidgets.QLabel('Average length of turn:')
        labelTurnMin = QtWidgets.QLabel('Average number of turns:')
        idLayout.addWidget(labelID)
        idLayout.addWidget(labelRST)
        idLayout.addWidget(labelTurn)
        idLayout.addWidget(labelTurnLen)
        idLayout.addWidget(labelTurnMin)

        self.measuresLayout.insertLayout(-1, idLayout)

        rstLabels = []
        turnLabels = []
        turnLenLabels = []
        turnMinLabels = []
        for id_no in range(numberIDs):
            idLayout = QtWidgets.QVBoxLayout()
            color = tuple([int(a * 255) for a in colors[id_no]])

            labelID = QtWidgets.QLabel('ID%i' % id_no)
            labelID.setFixedWidth(60)
            labelID.setFixedHeight(20)
            labelID.setStyleSheet('font: bold 12px; color: black; background-color: rgb(%i,%i,%i)' % color)

            labelRST = QtWidgets.QLabel('%.2f' % speak[id_no][1])
            rstLabels.append(labelRST)
            labelTurn = QtWidgets.QLabel('%i' % speak[id_no][2])
            turnLabels.append(labelTurn)
            labelTurnLen = QtWidgets.QLabel('%.2f sec' % speak[id_no][3])
            turnLenLabels.append(labelTurnLen)
            labelTurnMin = QtWidgets.QLabel('%.2f /min' % speak[id_no][4])
            turnMinLabels.append(labelTurnMin)

            idLayout.addWidget(labelID)
            idLayout.addWidget(labelRST)
            idLayout.addWidget(labelTurn)
            idLayout.addWidget(labelTurnLen)
            idLayout.addWidget(labelTurnMin)

            self.measuresLayout.insertLayout(-1, idLayout)

        self.labels['RST'] = rstLabels
        self.labels['Turns'] = turnLabels
        self.labels['TurnLength'] = turnLenLabels
        self.labels['TurnPerMin'] = turnMinLabels

        self.measures.setLayout(self.measuresLayout)

    @QtCore.pyqtSlot(list, int)
    def initSpeakingGraph(self, colors, numberIDs):
        """ initialize speaking graph """

        self.speakBlob = pg.ScatterPlotItem([0], [0], size=0.3, pxMode=False,
                           pen=pg.mkPen(color=0.5, width=2), brush=pg.mkBrush(255, 255, 255))
        self.speakBlob.setZValue(100)
        self.speakingGraph.addItem(self.speakBlob)

        pos = 360 // numberIDs
        for id_no in range(numberIDs):
            color = tuple([int(a * 255) for a in colors[id_no]])
            idPos = pos * id_no
            x = 2 * np.cos(np.deg2rad(idPos))
            y = 2 * np.sin(np.deg2rad(idPos))
            scatterPlot = pg.ScatterPlotItem([x], [y], size=0.2, pxMode=False,
                                             pen=pg.mkPen(color=color, width=2), brush=pg.mkBrush(*color))
            self.speakingGraph.addItem(scatterPlot)
            self.positions[id_no] = [x, y]
            self.idBlobs[id_no] = scatterPlot

            plt = self.speakingGraph.plot(x=[], y=[], pen=pg.mkPen(color=color, width=1))
            self.speakingGraphPlots.append(plt)

    @QtCore.pyqtSlot(dict, int, int)
    def updateSpeakingGraph(self, rst, active_speaker, numberIDs):
        """ frame updates for gaze graph """
        blobPos = np.array([0, 0])
        for id_no in range(numberIDs):
            diff = rst[id_no] - (1.0 / numberIDs)
            if diff > 0:
                blobPos = blobPos + 2 * diff * np.array(self.positions[id_no])
        self.speakBlob.setData([blobPos[0]], [blobPos[1]])

        for id_no in range(numberIDs):
            color = tuple([int(a * 255) for a in self.colors[id_no]])
            pen = pg.mkPen(color=color, width=(rst[id_no]*4))

            x = [self.positions[id_no][0], blobPos[0]]
            y = [self.positions[id_no][1], blobPos[1]]
            self.speakingGraphPlots[id_no].setData(x, y)
            self.speakingGraphPlots[id_no].setPen(pen)

            if id_no == active_speaker:
                self.idBlobs[id_no].setData([self.positions[id_no][0]], [self.positions[id_no][1]],
                                            pen=pg.mkPen(color=(0, 0, 0), width=2))
            else:
                self.idBlobs[id_no].setData([self.positions[id_no][0]], [self.positions[id_no][1]],
                                            pen=pg.mkPen(color=color, width=2))