vispy/vis.py at master · PioSte/vispy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import matplotlib.pyplot as plt


class IndexTracker:
    def __init__(self, X, sliced_axis=2, axis_labels=['y','x','z']):
        self.fig, self.ax = plt.subplots(1, 1)

        self.axis_labels = axis_labels
        self.current_axis = int(sliced_axis)
        self.current_label = axis_labels[self.current_axis]

        self.X = X
        self.current_axis_size = self.X.shape[self.current_axis]
        # self.X_rot = self.X.copy()
        # rows, cols, self.slices = X.shape
        self.ind = self.X.shape[self.current_axis]//2

        self.y_mouse_prev = None

        self.im = self.ax.imshow(self.X[:, :, self.ind])
        # self.im = self.ax.imshow(self.X_rot[:, :, self.ind])
        self.fig.colorbar(self.im, ax=self.ax)

        self.update()
        self.ax.set_title('Mouse drag up/down to select slice')

    def onmousemove(self, event):
        if event.x is not None and event.y is not None:
            if str(event.button) == 'MouseButton.LEFT':
                steps = int(round((event.y - self.y_mouse_prev) / 2))
                self.ind = (self.ind + steps) % self.current_axis_size
                self.update()
            self.y_mouse_prev = event.y
        else:
            self.y_mouse_prev = None

    def onrightmouseclick(self, event):
        if str(event.button) == 'MouseButton.RIGHT':
            self.current_axis = int(self.current_axis + 1)
            self.current_axis = self.current_axis % 3
            self.current_label = self.axis_labels[self.current_axis]
            self.current_axis_size = self.X.shape[self.current_axis]
            self.ind = self.ind % self.current_axis_size
            self.ax.clear()
            if self.current_axis == 0:
                self.im = self.ax.imshow(self.X[self.ind])
            elif self.current_axis == 1:
                self.im = self.ax.imshow(self.X[:, self.ind, :])
            else:
                self.im = self.ax.imshow(self.X[..., self.ind])
            self.update()

    def update(self):
        if self.current_axis == 0:
            self.im.set_data(self.X[self.ind])
        elif self.current_axis == 1:
            self.im.set_data(self.X[:, self.ind, :])
        else:
            self.im.set_data(self.X[..., self.ind])
        self.ax.set_title(self.current_label + ' = ' + str(self.ind) + ' / ' + str(self.current_axis_size))
        self.im.axes.figure.canvas.draw()


def vis(vol, sliced_axis=2, axis_labels=['y','x','z']):
    tracker = IndexTracker(vol, sliced_axis=2, axis_labels=['y', 'x', 'z'])
    tracker.fig.canvas.mpl_connect('motion_notify_event', tracker.onmousemove)
    tracker.fig.canvas.mpl_connect('button_press_event', tracker.onrightmouseclick)
    plt.show(block=True)


if __name__ == "__main__":
    import numpy as np
    sampling = 100
    x = np.linspace(-1, 1, sampling)
    y = x.copy() + 0.1
    z = x.copy() - 0.2

    x = x[..., np.newaxis, np.newaxis]
    y = y[np.newaxis, ..., np.newaxis]
    z = z[np.newaxis, np.newaxis, ...]

    # x, y, z = np.meshgrid(x, y, z)
    r = np.sqrt(x**2 + y**2 + z**2)

    # ball = np.zeros_like(x)
    ball = np.zeros((sampling, sampling, sampling))
    ball[r < 0.8] = 1

    vis(ball)
    test = 1