~mrol-dev/mrol/trunk

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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
'''
3D visualiser
Author: Julian Ryde and Nick Hillier
'''

#Copyright 2010-2011, Julian Ryde and Nicholas Hillier.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU Lesser General Public License for more details.
#
#You should have received a copy of the GNU Lesser General Public License
#along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import division 

import visual
import numpy as np
import pylab
import time
import mrol_mapping.visualiser.dispxyz as dispxyz
from threading import Thread

showlocal = False

def _draw_axes(position, size=1):
    '''Draws a small set of x,y,z axes coloured red, green, blue and the 
    position specified.'''
    # TODO make it account for orientation and take a pose
    position = np.array(position)
    visual.curve(pos=(position, position + (size,0,0)), color=(1,0,0))
    visual.curve(pos=(position, position + (0,size,0)), color=(0,1,0))
    visual.curve(pos=(position, position + (0,0,size)), color=(0,0,1))

class RotateVisual(Thread):
    def __init__(self, scene):
        Thread.__init__(self)
        self.scene_angle = 0
        self.scene = scene

    def run(self):
        while True:
            self.scene_angle += 0.005
            self.scene.forward = (np.cos(self.scene_angle), np.sin(self.scene_angle), -0.5) 
            self.scene.up = (0, 0, 1)
            #self.scene.center += (0.1, 0, 0)
            time.sleep(0.01)

class Visualiser:
    def __init__(self, npts=100000, Rotate=False, title=''):
        ptsize = 3
        mapptsize = 2
        self.scene_angle = 0
        if showlocal:
            self.localscene = visual.display(title=title + ' Local Coordinate System')
        # generate new window
        #self.globalscene = visual.display()
        # reuse existing scene window
        self.globalscene = visual.display.get_selected()

        # cannot seem to set title on active window
        #self.globalscene.title = title + ' Global Coordinate System'
        self.globalscene.background = (1, 1, 1)
        #self.globalscene.up = (0, 0, 1)
        #self.globalscene.forward = (0, 1, -0.4)
        self.currentrow = 0

        self.auto_colour = False

        # Set up the global coordinate frame visualiser
        self.globalscene.select()
        # initial view is looking down from top
        #self.globalscene.forward = (0,1.0,0)
        #self.globalscene.up =(0, 0, 1.0)
        #self.globalscene.forward =(0, 0, -1.0)

        #self.visualrobot = visual.box(length=3, height=2, width=1.5)
        w = 1
        wid = 0.2
        self.robotpts = np.array(( (0, -wid, 0, w), (0, wid, 0, w), (3*wid, 0, 0, w), (0, -wid, 0, w) ))
        #self.visualrobot = visual.points(pos=self.robotpts[:, :3], size=4, shape='round', color=(0, 1, 1))
        self.visualrobot = visual.curve(pos=self.robotpts[:, :3], color=(0, 1, 1))

        # draw the origin coordinate system
        _draw_axes((0,0,0))



        X = np.array([(-1, -1), (-1, 1), (1, 1), (1, -1), (-1, -1)])
        square = visual.curve(pos=50*X)
        self.leftmappts = visual.points(size=ptsize, shape='square', color=(1, 0, 0))
        self.rightmappts = visual.points(size=ptsize, shape='square', color=(0, 1, 0))
        self.spinmappts = visual.points(size=ptsize, shape='square', color=(0, 0, 1))
        X = np.zeros((npts, 3))
        self.mappts = visual.points(pos=X, size=mapptsize, shape='square', color=(1, 1, 1))
        self.trajpts_ind = 0
        self.trajpts = visual.points(pos=np.zeros((10000, 3)), color=(0, 0, 1), size=3)
        visual.scene.show_rendertime = True

        if Rotate:
            # Enable continuous rotation
            RV = RotateVisual(self.globalscene)
            RV.start()
        else:
            # Enable mouse panning
            MT = dispxyz.EnableMouseThread(self.globalscene)
            MT.start()

        # Set up the local coordinate frame visualiser
        if showlocal:
            self.localscene.select()
            self.leftpts = visual.points(color=(1, 0, 0), size=ptsize, shape='square')
            self.rightpts = visual.points(color=(0, 1, 0), size=ptsize, shape='square')
            self.spinpts = visual.points(color=(0, 0, 1), size=ptsize, shape='square')
            visual.scene.show_rendertime = True
        
        self.colourmin = -2
        self.colourmax = 2

    def getkeypress(self):
        key = None
        if self.globalscene.kb.keys:
            key = self.globalscene.kb.getkey()
        return key
    
    def setminmax(self,colourmin,colourmax):
        self.colourmin = colourmin
        self.colourmax = colourmax
    
    def setautominmax(self, Bool=True):
        self.auto_colour = Bool

    def set_orthographic(self, Bool=False):
        if Bool:
            self.globalscene.fov = 1e-5

    def setrobotpose(self, transform_mat): # TODO is there are more elegant way of doing this?
        self.visualrobot.pos[:] = np.dot(transform_mat, self.robotpts.T).T[:, :3]

    # TODO refactor this left and right
    def setleftpts(self, xyzs, M=None):
        if showlocal: self.leftpts.pos = xyzs[:, :3]
        if M != None:
            xyzs = np.dot(M, xyzs.T).T
        self.leftmappts.pos = xyzs[:, :3]
        #for i in X[:, :3]:
            #self.leftmappts.append(i)

    def setrightpts(self, xyzs, M=None):
        if showlocal: self.rightpts.pos = xyzs[:, :3]
        if M != None:
            xyzs = np.dot(M, xyzs.T).T
        self.rightmappts.pos = xyzs[:, :3]
        #for i in X[:, :3]:
            #self.rightmappts.append(i)

    def setspinpts(self, xyzs, M=None):
        if showlocal: self.spinpts.pos = xyzs[:, :3]
        if M != None:
            xyzs = np.dot(M, xyzs.T).T
        self.spinmappts.pos = xyzs[:, :3]

    # TODO make sure the ids of the pos arrays do not change to avoid memory leaks

    def clear(self):
        self.mappts.pos[:] = 0
        self.currentrow = 0

    def addmappts(self, xyzs, colour=None):
        # add this to the array of points to visualise
        # TODO make this more elegant so that you do not waste space at the end and 
        # do not leave points hanging around?
        # something like
        # X[:, 0].put([9, 10], [99, 98], mode='wrap') for the X, Y and Z coords

        # TODO fix bug if xyzs is longer than self.npts
        assert len(xyzs) < len(self.mappts.pos), 'Adding more points than visualiser was initialised with'
        # handle cases when xzys is empty
        if len(xyzs) < 1: 
            return
        if self.currentrow+len(xyzs) > len(self.mappts.pos): 
            # if adding the points would go off the end of the array start 
            # adding at the beginning again
            self.currentrow = 0 
        cr = self.currentrow
        self.mappts.pos[cr:cr+len(xyzs)] = xyzs[:, :3]
        if colour == None:
            X = xyzs[:, 2]
            if self.auto_colour == True:
                self.colourmin = np.min(X)
                self.colourmax = np.max(X)
            colourmin = self.colourmax#-4
            colourmax = self.colourmin#4
            colours = (X-colourmin)/(colourmax-colourmin)
            self.mappts.color[cr:cr+len(xyzs)] = pylab.cm.gist_rainbow(colours)[:, :3]
            # jet, Accent, Blues, copper, autumn, cool, gist_earth, gist_heat, 
            # gist_rainbow, spectral, spring, summer, winter
        else:
            self.mappts.color[cr:cr+len(xyzs)] = colour
        self.currentrow += len(xyzs)

    def setinfostr(self, labeltext):
        self.infolabel = visual.label()
        self.infolabel.pos = (-20, 0, 20)
        # TODO is this a memory leak?
        self.infolabel.text = str(labeltext)

    # TODO make a circular buffer type array as this is often needed by this 
    # visualiser code
    def addtrajectorypoint(self, xyz):
        self.trajpts.pos[self.trajpts_ind, :] = xyz
        self.trajpts_ind += 1
        _draw_axes(xyz, size=0.1)
        if self.trajpts_ind >= len(self.trajpts.pos):
            self.trajpts_ind = 0