~wecacuee/mrol/mrol-dev

« back to all changes in this revision

Viewing changes to mrol_mapping/visualiser/robotvisualiser.py

  • Committer: Julian Ryde
  • Date: 2011-06-03 16:52:15 UTC
  • Revision ID: julian_ryde-20110603165215-oqr6yo9nwng3j9nb
Initial import from revision 206 of Julian Ryde's mrol_mapping repository.

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
from __future__ import division 
 
2
'''Autonomous System 3D visualiser, from data file or live from the store
 
3
 
 
4
TODO currently reads from log file only need to add ability to read from store
 
5
 
 
6
Author: Julian Ryde'''
 
7
 
 
8
import visual
 
9
import numpy as np
 
10
import pylab
 
11
import time
 
12
import mrol_mapping.visualiser.dispxyz as dispxyz
 
13
from threading import Thread
 
14
 
 
15
showlocal = False
 
16
 
 
17
class RotateVisual(Thread):
 
18
    def __init__(self, scene):
 
19
        Thread.__init__(self)
 
20
        self.scene_angle = 0
 
21
        self.scene = scene
 
22
 
 
23
    def run(self):
 
24
        while True:
 
25
            self.scene_angle += 0.005
 
26
            self.scene.forward = (np.cos(self.scene_angle), np.sin(self.scene_angle), -0.5) 
 
27
            self.scene.up = (0, 0, 1)
 
28
            #self.scene.center += (0.1, 0, 0)
 
29
            time.sleep(0.01)
 
30
 
 
31
class Visualiser:
 
32
    def __init__(self, npts=100000, Rotate=False, title=''):
 
33
        ptsize = 3
 
34
        mapptsize = 1
 
35
        self.scene_angle = 0
 
36
        self.localscene = visual.display(title=title + ' Local Coordinate System')
 
37
        self.globalscene = visual.display(title=title + ' Global Coordinate System')
 
38
        #self.globalscene.background = (1, 1, 1)
 
39
        #self.globalscene.up = (0, 0, 1)
 
40
        #self.globalscene.forward = (0, 1, -0.4)
 
41
        self.currentrow = 0
 
42
 
 
43
        self.auto_colour = False
 
44
 
 
45
        # Set up the global coordinate frame visualiser
 
46
        self.globalscene.select()
 
47
 
 
48
        self.globalscene.up =(0, 0, 1.0)
 
49
 
 
50
        #self.visualrobot = visual.box(length=3, height=2, width=1.5)
 
51
        w = 1
 
52
        wid = 0.2
 
53
        self.robotpts = np.array(( (0, -wid, 0, w), (0, wid, 0, w), (3*wid, 0, 0, w), (0, -wid, 0, w) ))
 
54
        #self.visualrobot = visual.points(pos=self.robotpts[:, :3], size=4, shape='round', color=(0, 1, 1))
 
55
        self.visualrobot = visual.curve(pos=self.robotpts[:, :3], color=(0, 1, 1))
 
56
 
 
57
        X = np.array([(-1, -1), (-1, 1), (1, 1), (1, -1), (-1, -1)])
 
58
        square = visual.curve(pos=50*X)
 
59
        self.leftmappts = visual.points(size=ptsize, shape='square', color=(1, 0, 0))
 
60
        self.rightmappts = visual.points(size=ptsize, shape='square', color=(0, 1, 0))
 
61
        self.spinmappts = visual.points(size=ptsize, shape='square', color=(0, 0, 1))
 
62
        X = np.zeros((npts, 3))
 
63
        self.mappts = visual.points(pos=X, size=mapptsize, shape='square', color=(1, 1, 1))
 
64
        self.trajpts_ind = 0
 
65
        self.trajpts = visual.points(pos=np.zeros((10000, 3)), color=(0, 0, 1), size=2)
 
66
        visual.scene.show_rendertime = True
 
67
 
 
68
        if Rotate:
 
69
            # Enable continuous rotation
 
70
            RV = RotateVisual(self.globalscene)
 
71
            RV.start()
 
72
        else:
 
73
            # Enable mouse panning
 
74
            MT = dispxyz.EnableMouseThread(self.globalscene)
 
75
            MT.start()
 
76
 
 
77
        # Set up the local coordinate frame visualiser
 
78
        if showlocal:
 
79
            self.localscene.select()
 
80
            self.leftpts = visual.points(color=(1, 0, 0), size=ptsize, shape='square')
 
81
            self.rightpts = visual.points(color=(0, 1, 0), size=ptsize, shape='square')
 
82
            self.spinpts = visual.points(color=(0, 0, 1), size=ptsize, shape='square')
 
83
            visual.scene.show_rendertime = True
 
84
        
 
85
        self.colourmin = -4
 
86
        self.colourmax = 4
 
87
 
 
88
    def getkeypress(self):
 
89
        key = None
 
90
        if self.globalscene.kb.keys:
 
91
            key = self.globalscene.kb.getkey()
 
92
        return key
 
93
    
 
94
    def setminmax(self,colourmin,colourmax):
 
95
        self.colourmin = colourmin
 
96
        self.colourmax = colourmax
 
97
    
 
98
    def setautominmax(self, Bool=True):
 
99
        self.auto_colour = Bool
 
100
 
 
101
    def set_orthographic(self, Bool=False):
 
102
        if Bool:
 
103
            self.globalscene.fov = 1e-5
 
104
 
 
105
    def setrobotpose(self, transform_mat): # TODO is there are more elegant way of doing this?
 
106
        self.visualrobot.pos[:] = np.dot(transform_mat, self.robotpts.T).T[:, :3]
 
107
 
 
108
    # TODO refactor this left and right
 
109
    def setleftpts(self, xyzs, M=None):
 
110
        if showlocal: self.leftpts.pos = xyzs[:, :3]
 
111
        if M != None:
 
112
            xyzs = np.dot(M, xyzs.T).T
 
113
        self.leftmappts.pos = xyzs[:, :3]
 
114
        #for i in X[:, :3]:
 
115
            #self.leftmappts.append(i)
 
116
 
 
117
    def setrightpts(self, xyzs, M=None):
 
118
        if showlocal: self.rightpts.pos = xyzs[:, :3]
 
119
        if M != None:
 
120
            xyzs = np.dot(M, xyzs.T).T
 
121
        self.rightmappts.pos = xyzs[:, :3]
 
122
        #for i in X[:, :3]:
 
123
            #self.rightmappts.append(i)
 
124
 
 
125
    def setspinpts(self, xyzs, M=None):
 
126
        if showlocal: self.spinpts.pos = xyzs[:, :3]
 
127
        if M != None:
 
128
            xyzs = np.dot(M, xyzs.T).T
 
129
        self.spinmappts.pos = xyzs[:, :3]
 
130
 
 
131
    # TODO make sure the ids of the pos arrays do not change to avoid memory leaks
 
132
 
 
133
    def clear(self):
 
134
        self.mappts.pos[:] = 0
 
135
        self.currentrow = 0
 
136
 
 
137
    def addmappts(self, xyzs, colour=None):
 
138
        # add this to the array of points to visualise
 
139
        # TODO make this more elegant so that you do not waste space at the end and 
 
140
        # do not leave points hanging around?
 
141
        # something like
 
142
        # X[:, 0].put([9, 10], [99, 98], mode='wrap') for the X, Y and Z coords
 
143
 
 
144
        # TODO fix bug if xyzs is longer than self.npts
 
145
        assert len(xyzs) < len(self.mappts.pos), 'Adding more points than visualiser was initialised with'
 
146
        # handle cases when xzys is empty
 
147
        if len(xyzs) < 1: 
 
148
            return
 
149
        if self.currentrow+len(xyzs) > len(self.mappts.pos): 
 
150
            # if adding the points would go off the end of the array start 
 
151
            # adding at the beginning again
 
152
            self.currentrow = 0 
 
153
        cr = self.currentrow
 
154
        self.mappts.pos[cr:cr+len(xyzs)] = xyzs[:, :3]
 
155
        if colour == None:
 
156
            X = xyzs[:, 2]
 
157
            if self.auto_colour == True:
 
158
                self.colourmin = np.min(X)
 
159
                self.colourmax = np.max(X)
 
160
            colourmin = self.colourmin#-4
 
161
            colourmax = self.colourmax#4
 
162
            colours = (X-colourmin)/(colourmax-colourmin)
 
163
            self.mappts.color[cr:cr+len(xyzs)] = pylab.cm.gist_rainbow(colours)[:, :3]
 
164
            # jet, Accent, Blues, copper, autumn, cool, gist_earth, gist_heat, 
 
165
            # gist_rainbow, spectral, spring, summer, winter
 
166
        else:
 
167
            self.mappts.color[cr:cr+len(xyzs)] = colour
 
168
        self.currentrow += len(xyzs)
 
169
 
 
170
    def setinfostr(self, labeltext):
 
171
        self.infolabel = visual.label()
 
172
        self.infolabel.pos = (-20, 0, 20)
 
173
        # TODO is this a memory leak?
 
174
        self.infolabel.text = str(labeltext)
 
175
 
 
176
    # TODO make a circular buffer type array as this is often needed by this 
 
177
    # visualiser code
 
178
    def addtrajectorypoint(self, xyz):
 
179
        self.trajpts.pos[self.trajpts_ind, :] = xyz
 
180
        self.trajpts_ind += 1
 
181
        if self.trajpts_ind >= len(self.trajpts.pos):
 
182
            self.trajpts_ind = 0