← Back to branch summary

~mrol-dev/mrol/trunk

« back to all changes in this revision

Viewing changes to rosutils_mrol/rosutils_mrol.py

  • Committer: Vikas Dhiman
  • Date: 2012-05-14 16:28:10 UTC
  • Revision ID: wecacuee@gmail.com-20120514162810-gqqevqvisr4qyi5r
Backtracing must be working fine now.

Show diffs side-by-side

added added

removed removed

Lines of Context:
rosutils_mrol/rosutils_mrol.py
1
 
'''Author: Nick Hillier'''
2
 
# Various utils for interacting with ROS
3
 
 
4
 
import rospy
5
 
 
6
 
import numpy as np
7
 
import tf
8
 
import mrol_mapping.poseutil as poseutil
9
 
import mrol_mapping.depth_image as depth_image
10
 
import cv
11
 
from cv_bridge import CvBridge, CvBridgeError
12
 
 
13
 
def makePose3D(transform):
14
 
    p = poseutil.Pose3D()
15
 
    R = tf.transformations.quaternion_matrix(transform[1])
16
 
    T = transform[0]
17
 
    #M = [R,T;0,0,0,0]
18
 
    M = np.zeros([4,4])
19
 
    M = R
20
 
    M[0:3,3:4] = np.asarray(T).reshape(3,1)    
21
 
    p.setMatrix(M)
22
 
    return p
23
 
 
24
 
from sensor_msgs.msg import PointCloud2
25
 
from sensor_msgs.msg import PointField 
26
 
def xyzs2pc2(xyzs,stamp,seq,frame_id):
27
 
    #import pydb;pydb.set_trace()
28
 
    pc2Data = PointCloud2()
29
 
    if len(xyzs) != 0:
30
 
        if type(xyzs[0][0]) == np.float64:
31
 
            xyzs = np.array(xyzs,dtype=np.float32)
32
 
        assert (xyzs.shape[1] != 3) or (xyzs.shape[1] != 6),\
33
 
                "Shape: {0}".format(xyzs.shape)
34
 
        if xyzs.shape[1] == 3:
35
 
            numfields = 3
36
 
        else:
37
 
            numfields = 4
38
 
        float32LE = np.dtype(np.float32,align='<')
39
 
        pc2Data.is_dense = False
40
 
        pc2Data.is_bigendian = False
41
 
        pc2Data.height = 1
42
 
        pc2Data.width = len(xyzs)
43
 
        pc2Data.header.frame_id = frame_id
44
 
        pc2Data.header.seq = seq
45
 
        pc2Data.header.stamp = stamp;
46
 
        for i in range(numfields):
47
 
            pc2Field = PointField()
48
 
            if i == 0:
49
 
                pc2Field.name = 'x'
50
 
            elif i == 1:
51
 
                pc2Field.name = 'y'
52
 
            elif i == 2:
53
 
                pc2Field.name = 'z'
54
 
            elif i == 3:
55
 
                pc2Field.name = 'rgb'
56
 
            else:
57
 
                pc2Field.name = 'unknown'
58
 
            pc2Field.count = 1
59
 
            #if type(xyzs[0][0]) == np.float64:
60
 
            #    pc2Field.datatype = 8 # FLOAT64
61
 
            #    dataLen = 8 # FLOAT64
62
 
            if type(xyzs[0][0]) == float32LE:
63
 
                pc2Field.datatype = 7 # FLOAT32
64
 
                dataLen = 4 # FLOAT32
65
 
            else:
66
 
                assert False, "unknown datatype"
67
 
            
68
 
            pc2Field.offset = i*dataLen
69
 
            pc2Data.fields.append(pc2Field)
70
 
        pc2Data.point_step = i*dataLen + dataLen;
71
 
        pc2Data.row_step = pc2Data.point_step * pc2Data.width;
72
 
        if xyzs.shape[1] == 3:
73
 
            data_buf = xyzs.flatten().tostring()
74
 
        else:
75
 
            #import pydb;pydb.set_trace()
76
 
            rgb = np.hstack([np.uint8(xyzs[:,6:2:-1]),
77
 
                             np.zeros([xyzs.shape[0],1],dtype=np.uint8)]
78
 
                           ).copy()
79
 
            rgb = rgb.view(dtype=float32LE)
80
 
            data = np.hstack([xyzs[:,0:3],rgb])
81
 
            data_buf = data.flatten().tostring()
82
 
        if (xyzs.shape[1] == 3):
83
 
            assert len(data_buf)/float(dataLen)/3.0 == len(xyzs), "len(data_buf)/4.0/3.0 = "+np.str(len(data_buf)/4.0/3.0)+" but len(xyzs) = "+np.str(len(xyzs))
84
 
        else:
85
 
            assert len(data_buf)/float(dataLen)/4.0 == len(xyzs), "len(data_buf)/4.0/4.0 = "+np.str(len(data_buf)/4.0/4.0)+" but len(xyzs) = "+np.str(len(xyzs))
86
 
        pc2Data.data = data_buf
87
 
    return pc2Data
88
 
    
89
 
def pc22xyzs(msg, max_range=100.):
90
 
    #import pydb;pydb.set_trace()
91
 
    assert msg.is_bigendian == False, "Big Endian support is not yet implemented"
92
 
    num_fields = len(msg.fields)
93
 
    assert num_fields >=3, "Message must contain at least 3 dimensions (x,y,z)"
94
 
    # TODO: checks on fields content for type etc
95
 
    float32LE = np.dtype(np.float32,align='<')
96
 
    #data = np.frombuffer(msg.data, dtype=float32LE)
97
 
    #data_arr = data.reshape((-1,num_fields*2))
98
 
    data = np.frombuffer(msg.data, dtype=np.uint8)
99
 
    data_arr = data.reshape((-1,msg.point_step))
100
 
    # TODO: WARNING! The copying below assumes msg.fields is x,y,z(,rgb).
101
 
    #xyzsD = data_arr[:,0:3].copy()
102
 
    data_arrF = data_arr.view(dtype=float32LE)
103
 
    xyzsD = data_arrF[:,0:3]
104
 
    inds = np.logical_not(np.isnan(xyzsD[:,1]))
105
 
    xyzsD = xyzsD[inds,:]
106
 
    ranges = (xyzsD[:,0]*xyzsD[:,0] + xyzsD[:,1]*xyzsD[:,1]) + xyzsD[:,2]*xyzsD[:,2]
107
 
    indsD = np.logical_and(ranges < max_range**2, ranges > 0)
108
 
    xyzsD = xyzsD[indsD,:]
109
 
 
110
 
    if num_fields == 4:
111
 
        # TODO: This indexing ONLY WORKS FOR KINECT!!!
112
 
        #import pydb;pydb.set_trace()
113
 
        #xyzsRGB = data_arrF[:,4] 
114
 
        xyzsRGB = data_arr[:,18:15:-1] # RGB is backwards!
115
 
        xyzsRGB = xyzsRGB[inds,:]
116
 
        xyzsRGB = xyzsRGB[indsD,:]
117
 
        #xyzsRGB = xyzsRGB.view(np.uint8)
118
 
        #xyzsRGB = xyzsRGB.reshape((-1,4))
119
 
        
120
 
        xyzs = np.hstack([xyzsD,xyzsRGB])
121
 
    else:
122
 
        xyzs = xyzsD
123
 
    
124
 
    return xyzs
125
 
 
126
 
def imgmsg_to_cv(imgmsg):
127
 
    return CvBridge().imgmsg_to_cv(imgmsg)
128
 
 
129
 
def cv2array(im):
130
 
  depth2dtype = {
131
 
        cv.IPL_DEPTH_8U: 'uint8',
132
 
        cv.IPL_DEPTH_8S: 'int8',
133
 
        cv.IPL_DEPTH_16U: 'uint16',
134
 
        cv.IPL_DEPTH_16S: 'int16',
135
 
        cv.IPL_DEPTH_32S: 'int32',
136
 
        cv.IPL_DEPTH_32F: 'float32',
137
 
        cv.IPL_DEPTH_64F: 'float64',
138
 
    }
139
 
 
140
 
  arrdtype=im.depth
141
 
  a = np.fromstring(
142
 
         im.tostring(),
143
 
         dtype=depth2dtype[im.depth],
144
 
         count=im.width*im.height*im.nChannels)
145
 
  a.shape = (im.height,im.width,im.nChannels)
146
 
  return a
147
 
 
148
 
def imgmsg2array(imgmsg):
149
 
    img_or_mat = imgmsg_to_cv(imgmsg)
150
 
    if isinstance(img_or_mat, cv.cvmat):
151
 
        return np.asarray(img_or_mat)
152
 
    else:
153
 
        return cv2array(img_or_mat)
154
 
 
155
 
def imgmsg2xyz(imgmsg, **kwargs):
156
 
    arr = imgmsg2array(imgmsg)
157
 
    # image_raw topic is encoded as 16-bit integer where each pixel depth is
158
 
    # in milimeters.
159
 
    # http://ros.org/reps/rep-0118.html
160
 
    # 
161
 
    # Convert it into SI units.
162
 
    arr = np.asarray(arr, dtype=np.float32)
163
 
    arr /= 1000.0
164
 
    assert(len(arr) > 0)
165
 
    return depth_image.image_to_points(arr, **kwargs)