VisPy Context Menu Items

gistfile1.py

import sys
import numpy as np

import sip
sip.setapi('QString', 2)
from PyQt4 import QtGui, QtCore
from vispy import app, scene
from vispy.scene import visuals
from vispy.geometry import generation as gen, create_sphere

#app.use_app('pyside')

class DemoScene(QtGui.QWidget):
    def __init__(self, keys='interactive'):
        self.keymapping = {'z':self.setZ,
                           'x':self.setX,
                           'y':self.setY,
                           'Space':self.setZoom,
                           '=':self.largermarker,
                           '+':self.largermarker,
                           '-':self.smallermarker,
                           's':self.createsphere}
        super(DemoScene, self).__init__()
        
        #Layout and canvas creation
        box = QtGui.QVBoxLayout(self)
        self.resize(500,500)
        self.setLayout(box)
        self.canvas = scene.SceneCanvas(keys=keys)
        box.addWidget(self.canvas.native)

        #Connect events
        self.canvas.events.mouse_press.connect(self.on_mouse_press)
        self.canvas.events.mouse_release.connect(self.on_mouse_release)
        self.canvas.events.mouse_move.connect(self.on_mouse_move)

        #Setup some defaults
        self.selected = []
        self.white = (1.0, 1.0, 1.0, 1.0)
        self.black = (0.0, 0.0, 0.0, 0.0)

        #Selection
        self.rectselect = False
        self.selectionpolygon = visuals.Polygon(color=None,
                                                border_color=(1.0, 0.0, 0.0, 1.0))

        #Context Menu
        self.context = False

        #Sphere
        self.mdata = create_sphere(200, 200, 1.0)
        self.mesh = visuals.Mesh(meshdata=self.mdata,
                            face_colors=(1.0, 0.0, 0.0, 1.0),
                            color=(0.5, 0.5, 0.5, .75))

        #Camera
        self.view = self.canvas.central_widget.add_view()
        self.view.camera = scene.cameras.TurntableCamera(elevation = 25, azimuth=20, distance = 2.0, center=(0,0,0), width=20)
        self.view.camera.viewbox.clip_method = 'viewport'
        self.view.camera.set_ortho()
        self.view.camera._update_camera_pos()

        #Test data
        self.data = np.random.uniform(-1, 1, size=(200, 3))
        self.facecolor = np.ones((200,4), dtype=np.float)
        self.ptsize = 3
        self.scatter = visuals.Markers()
        self.scatter.set_data(self.data, face_color=self.facecolor, edge_color=None, size=self.ptsize)

        #Add scatter to view and find the transform
        self.view.add(self.scatter)
        self.tr = self.scatter.node_transform(self.canvas.canvas_cs)

        # Add a 3D axis to keep us oriented
        axis = scene.visuals.XYZAxis(parent=self.view.scene)

    def markselected(self):
        """
        Change the color of the selected point
        """
        self.facecolor[self.facecolor[:,1] != 1.0] = self.white
        self.scatter.set_data(self.data, face_color=self.facecolor, size=self.ptsize)
        for i in self.selected:
            self.facecolor[i] = [1.0, 0.0, 0.0, 1]

        self.scatter.set_data(self.data, face_color = self.facecolor, size=self.ptsize)
        self.scatter.update()

    def keyPressEvent(self, event):
        """
        Bootstrap the Qt keypress event items
        """
        if event.text() in self.keymapping.keys():
            self.keymapping[event.text()]()
        elif event.key() == QtCore.Qt.Key_Shift:
            self.rectselect = True
        elif event.key() == QtCore.Qt.Key_Space:
            self.keymapping['Space']()
        elif event.key() == QtCore.Qt.Key_Alt:
            self.context = True
        else:
            pass
            #print event.text()

    def keyReleaseEvent(self, event):
        """
        Bootstrap the Qt Key release event
        """
        if event.key() == QtCore.Qt.Key_Shift:
            self.rectselect = False
        elif event.key() == QtCore.Qt.Key_Alt:
            self.context = False

    def on_mouse_press(self, event):
        """
        Mouse button press event
        """
        if event.button == 1 and self.rectselect == False:
            ##Ray intersection on the CPU to highlight the selected point(s)
            data = self.tr.simplified().map(self.data)[:,:2]
            m1 = data > (event.pos - 4)
            m2 = data < (event.pos + 4)
            self.selected = np.argwhere(m1[:,0] & m1[:,1] & m2[:,0] & m2[:,1])
            self.markselected()
        elif event.button == 2 and self.context == True:
            menu = QtGui.QMenu(self)
            if len(self.selected) == 0:
                menu.addAction('A')
                menu.addAction('B')
                menu.addAction('C')
            else:
                menu.addAction('D')
                menu.addAction('E')
                menu.addAction('F')

            action = menu.exec_(event.native.globalPos())

    def on_mouse_release(self, event):
        if event.button == 2 and self.context == False:
            #Selecting using a rectangle
            self.facecolor[self.facecolor[:,1] != 1.0] = self.white
            x1, y1 = event.last_event.pos
            x0, y0 = event.press_event.pos

            #Sort for the origin
            if x1 < x0:
                x0, x1 = x1, x0
            if y1 < y0:
                y0, y1 = y1, y0

            data = self.tr.simplified().map(self.data)[:,:2]
            m1 = data[:,0] < x1
            m2 = data[:,0] > x0
            m3 = data[:,1] < y1
            m4 = data[:,1] > y0

            m1 =m1 * m2
            m2 = m3 * m4

            self.selected = np.argwhere(m1 & m2)
            self.markselected()

            #Remove the selection polygon
            self.selectionpolygon.parents = []

    def on_mouse_move(self, event):
        if event.button == 2 and event.is_dragging and self.context == False:
            p1 = event.press_event.pos
            p2 = event.pos
            polydata = np.array([p1, [p1[0], p2[1]], p2, [p2[0], p1[1]]])
            self.selectionpolygon.parent = self.canvas.scene
            self.selectionpolygon.pos = polydata

    def setZ(self):
        self.view.camera.azimuth = 0
        self.view.camera.elevation = 90
        self.view.camera._update_camera_pos()

    def setX(self):
        self.view.camera.azimuth = 0
        self.view.camera.elevation = 0
        self.view.camera._update_camera_pos()

    def setY(self):
        self.view.camera.azimuth = 90
        self.view.camera.elevation = 0
        self.view.camera._update_camera_pos()

    def setZoom(self):
        #Need bounding box info / functionality
        pass

    def largermarker(self):
        self.ptsize += 2
        self.scatter._v_size_var = self.ptsize
        self.scatter._marker_fun['v_size'] = self.scatter._v_size_var
        self.scatter.update()

    def smallermarker(self):
        self.ptsize -= 2
        self.scatter._v_size_var = self.ptsize
        self.scatter._marker_fun['v_size'] = self.scatter._v_size_var
        self.scatter.update()

    def createsphere(self):
        if self.mesh.parents == ():
            self.view.add(self.mesh)
        else:
            self.mesh.parents = []
            self.canvas.update()


if __name__ == '__main__':
    appQt = QtGui.QApplication(sys.argv)
    view = DemoScene(keys='interactive')
    view.show()
    appQt.exec_()