int-brain-lab / iblrig / 12279337432

import json

import sys

import matplotlib.pyplot as plt

import numpy as np

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

from matplotlib.figure import Figure

from PyQt5.QtGui import QColor, QPalette

from PyQt5.QtWidgets import (

from ibllib.atlas import AllenAtlas, Insertion

from ibllib.tests import TEST_DB

from one.webclient import AlyxClient

ACTUAL_DB = {

def plot_trajectories(ax, names, trajectories, atlas=None):

    assert atlas

    top = atlas.top

    extent = np.hstack((atlas.bc.xlim, atlas.bc.ylim))

    ax.imshow(top, extent=extent, cmap='Greys_r')

    ax.set_xlim(atlas.bc.xlim)

    ax.set_ylim(atlas.bc.ylim)

    prop_cycle = plt.rcParams['axes.prop_cycle']

    colors = prop_cycle.by_key()['color']

    eps = 0.0001

    for name, traj, color in zip(names, trajectories, colors, strict=True):

        x = traj[0, 0]

        y = traj[0, 1]

        if x == y == 0:

            continue

        ax.plot(traj[:, 0], traj[:, 1])

        ax.plot([x], [y], 'o', color=color)

        ax.text(x - eps, y - 4 * eps, name, color=color)

class TrajectoryLoader:

    def __init__(self, atlas=None):

        self.alyx = AlyxClient(**TEST_DB)

        self.atlas = atlas

    def _save_rest(self, n, v='read', pk=None):

        d = self.alyx.rest(n, v, id=pk)

        with open(f'{n}.json', 'w') as f:

            json.dump(d, f, indent=1)

    def save_subject(self, pk):

        self._save_rest(self, 'subjects', pk=pk)

    def save_session(self, pk):

        self._save_rest(self, 'sessions', pk=pk)

    def save_insertion(self, pk):

        self._save_rest(self, 'insertions', pk=pk)

    def save_trajectories(self):

        self._save_rest(self, 'trajectories', v='list')

    def create(self, name, path):

        with open(path) as f:

            self.alyx.rest(name, 'create', data=json.load(f))

    def get_trajectory(self, chronic_insertion):

        trajectories = self.alyx.rest('trajectories', 'list', chronic_insertion=chronic_insertion)

        if not trajectories:

            return

        priorities = {

        trajectory = sorted(trajectories, key=lambda t: priorities.get(t['provenance'], 0))[-1]

        ins = Insertion.from_dict(trajectory, brain_atlas=self.atlas)

        return np.vstack((ins.entry, ins.tip))

    def get_trajectories(self, subject):

        chronic_insertions = self.alyx.rest('chronic-insertions', 'list', subject=subject, model='fiber')

        names = [i['name'] for i in chronic_insertions]

        trajectories = [self.get_trajectory(i['id']) for i in chronic_insertions]

        return names, trajectories

class MainWindow(QMainWindow):

    def __init__(self, nickname=None, names=None, trajectories=None):

        super().__init__()

        self.atlas = AllenAtlas(25)

        self.atlas.compute_surface()

        self.nickname = nickname

        self.names = names

        self.trajectories = trajectories

        self.setWindowTitle('Fiber insertions')

        main_widget = QWidget()

        main_layout = QVBoxLayout()

        top_panel = QWidget()

        top_layout = QVBoxLayout()

        label_subject = QLabel(self.nickname)

        top_layout.addWidget(label_subject)

        self.textboxes = []

        prop_cycle = plt.rcParams['axes.prop_cycle']

        colors = prop_cycle.by_key()['color']

        for i in range(len(self.trajectories)):

            color = colors[i]

            self.textboxes.append(QLineEdit())

            rl = QFormLayout()

            c = self.trajectories[i][0]  # 0 is entry point, 1 is tip

            s = f'{self.names[i]}: AP {c[0]:.4f}, ML {c[1]:.4f}, DV {c[2]:.4f}'

            label = QLabel(s)

            palette = label.palette()

            palette.setColor(QPalette.WindowText, QColor(color))

            label.setPalette(palette)

            rl.addRow(label, self.textboxes[i])

            top_layout.addLayout(rl)

        top_panel.setLayout(top_layout)

        bottom_panel = QWidget()

        bottom_layout = QVBoxLayout()

        self.figure = Figure()

        self.canvas = FigureCanvas(self.figure)

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

        bottom_layout.addWidget(self.canvas)

        bottom_panel.setLayout(bottom_layout)

        main_layout.addWidget(top_panel)

        main_layout.addWidget(bottom_panel)

        main_widget.setLayout(main_layout)

        self.setCentralWidget(main_widget)

        plot_trajectories(self.ax, self.names, self.trajectories, atlas=self.atlas)

if __name__ == '__main__':

    fig, ax = plt.subplots(1, 1)

    nickname = 'CQ004'

    names = ['NBM', 'PPT']

    trajectories = [

    app = QApplication(sys.argv)

    window = MainWindow(nickname, names, trajectories)

    window.show()

    sys.exit(app.exec_())