traja-gui.py

import os
from os.path import basename
from functools import partial
import sys

import matplotlib

matplotlib.use("Qt5Agg")
import pandas as pd
import matplotlib.pyplot as plt

plt.ioff()

import matplotlib.style as style
import matplotlib.backends.backend_qt5agg.FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import (
    FigureCanvasQTAgg as FigureCanvas,
    NavigationToolbar2QT as NavigationToolbar,
)
from PyQt5 import QtGui, QtWidgets, QtCore
from PyQt5.QtCore import Qt, QThread, QObject, pyqtSignal, pyqtSlot
from PyQt5.QtWidgets import QProgressBar, QMenu, QAction, QStatusBar

import traja

CUR_STYLE = "fast"
style.use(CUR_STYLE)
TIME_WINDOW = "30s"


class QtFileLoader(QObject):
    finished = pyqtSignal()
    progressMaximum = pyqtSignal(int)
    completed = pyqtSignal(list)
    intReady = pyqtSignal(int)

    def __init__(self, filepath):
        super(QtFileLoader, self).__init__()
        self.filepath = filepath

    @pyqtSlot()
    def read_in_chunks(self):
        """ load dataset in parts and update the progess par """
        chunksize = 10 ** 3
        lines_number = sum(1 for line in open(self.filepath))
        self.progressMaximum.emit(lines_number // chunksize)
        dfList = []

        # self.df = traja.read_file(
        #     str(filepath),
        #     index_col="time_stamps_vec",
        #     parse_dates=["time_stamps_vec"],
        # )

        TextFileReader = pd.read_csv(
            self.filepath,
            index_col="time_stamps_vec",
            parse_dates=["time_stamps_vec"],
            chunksize=chunksize,
        )
        for idx, df in enumerate(TextFileReader):
            df.index = pd.to_datetime(df.index, format="%Y-%m-%d %H:%M:%S:%f")
            dfList.append(df)
            self.intReady.emit(idx)
        self.completed.emit(dfList)
        self.finished.emit()


class PlottingWidget(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        # super(PrettyWidget, self).__init__()
        self.initUI()

    def initUI(self):
        self.setGeometry(600, 300, 1000, 600)
        self.center()
        self.setWindowTitle("Plot Trajectory")

        mainMenu = self.menuBar()
        fileMenu = mainMenu.addMenu("File")

        saveAction = QAction("Save as...")
        saveAction.setShortcut("Ctrl+S")
        saveAction.setStatusTip("Save plot to file")
        saveAction.setMenuRole(QAction.NoRole)
        saveAction.triggered.connect(self.file_save)
        fileMenu.addAction(saveAction)

        exitAction = QAction("&Exit", self)
        exitAction.setShortcut("Ctrl+Q")
        exitAction.setStatusTip("Exit Application")
        exitAction.setMenuRole(QAction.NoRole)
        exitAction.triggered.connect(self.close)
        fileMenu.addAction(exitAction)

        settingsMenu = mainMenu.addMenu("Settings")
        self.setStyleMenu = QMenu("Set Style", self)
        settingsMenu.addMenu(self.setStyleMenu)
        for style_name in ["default", "fast", "ggplot", "grayscale", "seaborn"]:
            styleAction = QAction(style_name, self, checkable=True)
            if style_name is CUR_STYLE:
                styleAction.setChecked(True)
            styleAction.triggered.connect(partial(self.set_style, style_name))
            self.setStyleMenu.addAction(styleAction)
        self.setTimeWindowMenu = QMenu("Set Time Window", self)
        settingsMenu.addMenu(self.setTimeWindowMenu)
        for window_str in ["None", "s", "30s", "H", "D"]:
            windowAction = QAction(window_str, self, checkable=True)
            if window_str is TIME_WINDOW:
                windowAction.setChecked(True)
            windowAction.triggered.connect(partial(self.set_time_window, window_str))
            self.setTimeWindowMenu.addAction(windowAction)

        # Grid Layout
        grid = QtWidgets.QGridLayout()
        widget = QtWidgets.QWidget(self)
        self.setCentralWidget(widget)
        widget.setLayout(grid)

        # Import CSV Button
        btn1 = QtWidgets.QPushButton("Import CSV", self)
        btn1.resize(btn1.sizeHint())
        btn1.clicked.connect(self.getCSV)
        grid.addWidget(btn1, 1, 0)

        # Canvas and Toolbar
        self.figure = plt.figure(figsize=(15, 5))
        self.canvas = FigureCanvas(self.figure)
        self.canvas.setContextMenuPolicy(Qt.CustomContextMenu)
        self.canvas.customContextMenuRequested.connect(self.popup)
        grid.addWidget(self.canvas, 2, 0, 1, 2)

        # DropDown mean / comboBox
        self.df = pd.DataFrame()
        self.columns = []
        self.plot_list = []

        self.comboBox = QtWidgets.QComboBox(self)
        self.comboBox.addItems(self.columns)
        grid.addWidget(self.comboBox, 0, 0)

        self.comboBox2 = QtWidgets.QComboBox(self)
        self.comboBox2.addItems(self.plot_list)
        grid.addWidget(self.comboBox2, 0, 1)

        # Plot Button
        btn2 = QtWidgets.QPushButton("Plot", self)
        btn2.resize(btn2.sizeHint())
        btn2.clicked.connect(self.plot)
        grid.addWidget(btn2, 1, 1)

        # Progress bar
        self.progress = QProgressBar(self)
        # self.progress.setRange(0, 1)
        grid.addWidget(self.progress, 3, 0, 1, 2)

        self.statusBar = QStatusBar()
        self.setStatusBar(self.statusBar)
        self.show()

    def set_style(self, style_name: str):
        global CUR_STYLE
        self.statusBar.showMessage(f"Style set to {style_name}")
        actions = self.setStyleMenu.actions()
        CUR_STYLE = style_name
        for action in actions:
            if action.text() == CUR_STYLE:
                # print(f"✓ {CUR_STYLE}")
                action.setChecked(True)
            else:
                action.setChecked(False)
        print(f"Style set to {CUR_STYLE}")

    def popup(self, pos):
        menu = QMenu()
        saveAction = menu.addAction("Save...")
        action = menu.exec_(self.canvas.viewport().mapToGlobal(pos))
        if action == saveAction:
            self.file_save()

    def file_save(self, target="figure"):
        name = QtGui.QFileDialog.getSaveFileName(self, "Save File")
        if target == "figure":
            self.figure.savefig(name)

    def update_progress_bar(self, i: int):
        self.progress.setValue(i)
        max = self.progress.maximum()
        self.statusBar.showMessage(f"Loading ... {100*i/max:.0f}%")

    def set_progress_bar_max(self, max: int):
        self.progress.setMaximum(max)

    def clear_progress_bar(self):
        self.progress.hide()
        self.statusBar.showMessage("Completed.")

    def getCSV(self):
        self.statusBar.showMessage("Loading CSV...")
        filepath, _ = QtWidgets.QFileDialog.getOpenFileName(
            self, "Open CSV", (QtCore.QDir.homePath()), "CSV (*.csv *.tsv)"
        )

        if filepath != "":
            self.filepath = filepath
            self.loaderThread = QThread()
            self.loaderWorker = QtFileLoader(filepath)
            self.loaderWorker.moveToThread(self.loaderThread)
            self.loaderThread.started.connect(self.loaderWorker.read_in_chunks)
            self.loaderWorker.intReady.connect(self.update_progress_bar)
            self.loaderWorker.progressMaximum.connect(self.set_progress_bar_max)
            # self.loaderWorker.read_in_chunks.connect(self.df)
            self.loaderWorker.completed.connect(self.list_to_df)
            self.loaderWorker.completed.connect(self.clear_progress_bar)
            self.loaderThread.finished.connect(self.loaderThread.quit)
            self.loaderThread.start()

    @pyqtSlot(list)
    def list_to_df(self, dfs: list):
        df = pd.concat(dfs)
        self.df = df
        self.columns = self.df.columns.tolist()
        self.plot_list = ["Actogram", "Polar Bar", "Polar Histogram", "Trajectory"]
        self.comboBox.clear()
        self.comboBox.addItems(self.columns)
        self.comboBox2.clear()
        self.comboBox2.addItems(self.plot_list)
        self.statusBar.clearMessage()

    def mousePressEvent(self, QMouseEvent):
        if QMouseEvent.button() == Qt.RightButton:

            print("Right Button Clicked")

    def load_project_structure(self, startpath, tree):
        """
        Load Project structure tree
        :param startpath:
        :param tree:
        :return:
        """
        from PyQt5.QtWidgets import QTreeWidgetItem
        from PyQt5.QtGui import QIcon

        for element in os.listdir(startpath):
            path_info = startpath + "/" + element
            parent_itm = QTreeWidgetItem(tree, [os.path.basename(element)])
            if os.path.isdir(path_info):
                self.load_project_structure(path_info, parent_itm)
                parent_itm.setIcon(0, QIcon("assets/folder.ico"))
            else:
                parent_itm.setIcon(0, QIcon("assets/file.ico"))

    def set_time_window(self, window: str):
        global TIME_WINDOW
        TIME_WINDOW = window
        self.statusBar.showMessage(f"Time window set to {window}")
        actions = self.setTimeWindowMenu.actions()
        for action in actions:
            if action.text() == TIME_WINDOW:
                action.setChecked(True)
            else:
                action.setChecked(False)
        print(f"Time window set to {window}")

    def plot(self):
        plt.clf()

        plot_kind = self.comboBox2.currentText()
        self.statusBar.showMessage(f"Plotting {plot_kind}")
        projection = (
            "polar" if plot_kind in ["Polar Bar", "Polar Histogram"] else "rectilinear"
        )

        ax = self.figure.add_subplot(111, projection=projection)

        title = f"{basename(self.filepath)}"

        # TODO: Move mapping to separate method
        if plot_kind == "Actogram":
            displacement = traja.trajectory.calc_displacement(self.df)
            if TIME_WINDOW != "None":
                displacement = displacement.rolling(TIME_WINDOW).mean()
                # from pyqtgraph.Qt import QtGui, QtCore
            traja.plotting.plot_actogram(displacement, ax=ax, interactive=False)
        elif plot_kind == "Trajectory":
            traja.plotting.plot(self.df, ax=ax, interactive=False)
        elif plot_kind == "Quiver":
            traja.plotting.plot_quiver(self.df, ax=ax, interactive=False)
        elif plot_kind == "Polar Bar":
            traja.plotting.polar_bar(self.df, ax=ax, title=title, interactive=False)
        elif plot_kind == "Polar Histogram":
            traja.plotting.polar_bar(
                self.df, ax=ax, title=title, overlap=False, interactive=False
            )
        plt.tight_layout()
        self.canvas.draw()
        self.statusBar.clearMessage()

    def center(self):
        qr = self.frameGeometry()
        cp = QtWidgets.QDesktopWidget().availableGeometry().center()
        qr.moveCenter(cp)
        self.move(qr.topLeft())


def main():
    app = QtWidgets.QApplication(sys.argv)
    w = PlottingWidget()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()