Decode additional fields in pointcloud2 msgs.

importRosbag/importRosbag.py

@@ -40,6 +40,7 @@
 from struct import error as structError
 from tqdm import tqdm
 
+import time
 # Local imports
 
 from .messageTypes.common import unpackHeader
@@ -59,6 +60,7 @@
 import logging
 
 def importTopic(topic, **kwargs):
+    time_start = time.time()
     msgs = topic['msgs']
     topicType = topic['type'].replace('/','_')
     if topicType == 'dvs_msgs_EventArray': topicDict = import_dvs_msgs_EventArray(msgs, **kwargs)
@@ -76,6 +78,8 @@ def importTopic(topic, **kwargs):
         return None
     if topicDict:
         topicDict['rosbagType'] = topic['type']
+        elapsed_time = time.time() - time_start
+        print("Tiempo transcurrido: {} segundos".format(elapsed_time))
     return topicDict
 
 def readFile(filePathOrName):
@@ -177,6 +181,7 @@ def rekeyConnsByTopic(connDict):
 
 def importRosbag(filePathOrName, **kwargs):
     global disable_bar
+    kwargs.setdefault('log', 'ERROR')
     disable_bar = kwargs.get('disable_bar')
     loglevel = kwargs.get('log')
     numeric_level = getattr(logging, loglevel.upper(), None)

importRosbag/messageTypes/common.py

@@ -43,6 +43,9 @@ def unpackRosUint32(data, ptr):
 def unpackRosUint8(data, ptr):
     return unpack('=B', data[ptr:ptr+1])[0], ptr+1
 
+def unpackRosUint16(data, ptr):
+    return unpack('=B', data[ptr:ptr+2])[0], ptr+2
+
 def unpackRosString(data, ptr):
     stringLen = unpack('=L', data[ptr:ptr+4])[0]
     ptr += 4
@@ -53,6 +56,14 @@ def unpackRosString(data, ptr):
     ptr += stringLen
     return outStr, ptr
 
+def unpackStringData(data, stringLen, ptr):
+    try:
+        outStr = data[ptr:ptr+stringLen].decode('utf-8')
+    except UnicodeDecodeError:
+        outStr = 'UnicodeDecodeError'
+    ptr += stringLen
+    return outStr, ptr
+
 def unpackRosFloat64Array(data, num, ptr):
     return np.frombuffer(data[ptr:ptr+num*8], dtype=np.float64), ptr+num*8 
 

importRosbag/messageTypes/sensor_msgs_PointCloud2.py

@@ -22,79 +22,135 @@
 This function imports the ros message type defined at:
 http://docs.ros.org/api/sensor_msgs/html/msg/PointCloud2.html
 
-For simplicity, we're currently directly unpacking the format that we are 
-encountering in the data, which is x,y,z,_,rgb,_,_,_ 
-each as 32-bit little-endian floats
+In this new version of the pointcloud2 extract, we can automated unpacking
+all the fields that are included in pointcloud2 message in the format that is required.
 """
 
-#%%
+# %%
 
 
 from tqdm import tqdm
 import numpy as np
+import time
+
+from .common import unpackRosString, unpackRosUint8, unpackRosUint32, unpackRosTimestamp, unpackRosFloat32, \
+    unpackRosUint16, unpackStringData
 
-from .common import unpackRosString, unpackRosUint8, unpackRosUint32, \
-                    unpackRosTimestamp
 
 def importTopic(msgs, **kwargs):
     '''
     ros message is defined here:
-        http://docs.ros.org/api/geometry_msgs/html/msg/PoseStamped.html
-    the result is are np arrays of float64 for:
-        rotQ (4 cols, quaternion)
-        angV (3 cols)
-        acc (3 cols)
-        mag (3 cols)
-        temp (1 cols) - but I'll probably ignore this to start with
+        http://docs.ros.org/api/sensor_msgs/html/msg/PointCloud2.html
+    the result is a variable dimension dictionary with
+        - timestamp
+        - point structure (include all field that declares in pointcloud2 msg)
     '''
-    #tempAll = np.zeros((sizeOfArray, 1), dtype=np.float64)
-    #for msg in tqdm(msgs, position=0, leave=True):
     disable_bar = kwargs.get('disable_bar')
-    tsByMessage = []
-    pointsByMessage = []
+    DTYPE_VAR = {
+        '1': '=i1',  # INT8
+        '2': '=u1',  # UINT8
+        '3': '=i2',  # INT16
+        '4': '=u2',  # UINT16
+        '5': '=i4',  # INT32
+        '6': '=u4',  # UINT32
+        '7': '=f4',  # FLOAT32
+        '8': '=f8',  # FLOAT64
+    }
+    dtype_header_1 = {
+        'names': ['sequence_id', 'timestamp_seconds', 'timestamp_nseconds',  # List of Strings
+                  'frame_len'],
+        'formats': ['=u4','=u4','=u4','=u4'],                                # List of Strings
+        'offsets': [0, 4, 8, 12],                                            # List of INT
+        'itemsize': 16                                                       # INT
+    }
+    dtype_header_2 = {
+        'names': ['height', 'width', 'fieldSize'],                          # List of Strings
+        'formats': [ '=u4', '=u4', '=u4'],                                  # List of Strings
+        'offsets': [0, 4, 8],                                               # List of INT
+        'itemsize': 12                                                      # INT
+    }
+    dtype_header_3 = {
+        'names': ['is_Bigendian', 'PointStep', 'RowStep', 'data_size'],     # List of Strings
+        'formats': [ '=u1', '=u4', '=u4','=u4'],                            # List of Strings
+        'offsets': [0, 1, 5, 9],                                            # List of INT
+        'itemsize': 13                                                      # INT
+    }
+    dtype_fields = {
+        'names': [],                                                        # List of Strings
+        'formats': [],                                                      # List of Strings
+        'offsets': [],                                                      # List of INT
+        'itemsize': 0,                                                      # INT
+        'count': []                                                         # List of INT
+    }
+
+    tsByMessage_2 = []
+    pointsByMessage_2 = []
+
     for msg in tqdm(msgs, disable=disable_bar):
-
         data = msg['data']
         ptr = 0
-        seq, ptr = unpackRosUint32(data, ptr)
-        ts, ptr = unpackRosTimestamp(data, ptr)
-        frame_id, ptr = unpackRosString(data, ptr)
-        height, ptr = unpackRosUint32(data, ptr)
-        width, ptr = unpackRosUint32(data, ptr) 
+        # Reading first header (with numpy)
+        header_1 = np.frombuffer(data[ptr:ptr + dtype_header_1['itemsize']], dtype=dtype_header_1)
+        timestamp = np.float64(header_1['timestamp_seconds'][0])+np.float64(header_1['timestamp_nseconds'][0])*0.000000001
+        ptr += dtype_header_1['itemsize']
+        # Unpacking frame_id with unknown dimension
+        frame_len = header_1['frame_len'][0].astype(int)
+        frame_id, ptr = unpackStringData(data, frame_len, ptr)
+        # Reading second header
+        header_2 = np.frombuffer(data[ptr:ptr + dtype_header_2['itemsize']], dtype=dtype_header_2)
+        ptr += dtype_header_2['itemsize']
 
-        if width > 0 and height > 0:
+        #Assign variables
+        height = header_2['height'][0]
+        width = header_2['width'][0]
+        fieldSize = header_2['fieldSize'][0]
 
-            arraySize, ptr = unpackRosUint32(data, ptr)
-            for element in range(arraySize):
-                # Move through the field definitions - we'll ignore these
-                # until we encounter a file that uses a different set
-                name, ptr = unpackRosString(data, ptr)
-                offset, ptr = unpackRosUint32(data, ptr)
-                datatype, ptr = unpackRosUint8(data, ptr)
-                count, ptr = unpackRosUint32(data, ptr)
-
-            isBigendian, ptr = unpackRosUint8(data, ptr)
-            pointStep, ptr = unpackRosUint32(data, ptr)
-            rowStep, ptr = unpackRosUint32(data, ptr)
+        #Reading fields of pointcloud2
+        if width > 0 and height > 0:
+            for element in range(fieldSize):
+                names, ptr = unpackRosString(data, ptr)
+                offset = np.frombuffer(data[ptr:ptr+4], dtype='=u4')
+                ptr += 4
+                datatype = np.frombuffer(data[ptr:ptr+1], dtype='=u1')
+                ptr += 1
+                count = np.frombuffer(data[ptr:ptr + 4], dtype='=u4')
+                ptr += 4
+                #APPEND TO THE DICTIONARY ('dtype_fields')
+                dtype_fields['names'].append(names)
+                dtype_fields['offsets'].append(offset[0])
+                dtype_fields['formats'].append(str(DTYPE_VAR[str(datatype[0])]))
+                dtype_fields['count'].append(count[0])
+            # Reading third header
+            header_3 = np.frombuffer(data[ptr:ptr + dtype_header_3['itemsize']], dtype=dtype_header_3)
+            ptr += dtype_header_3['itemsize']
 
+            #Assign variables
+            PointStep = header_3['PointStep'][0]
+            dtype_fields['itemsize'] = PointStep
             numPoints = width * height
-            points = np.empty((numPoints, 3), dtype=np.float32)
-            arraySize, ptr = unpackRosUint32(data, ptr)
-            # assert arraySize = width*height
+
+            #Read pointcloud2 data with datatype obtained before
+            data_array = np.empty((numPoints, 1), dtype=dtype_fields)
             for x in range(width):
-                for y in range(height):            
-                    points[x*height + y, :] = np.frombuffer(data[ptr:ptr+12], dtype=np.float32)
-                    ptr += pointStep
-            pointsByMessage.append(points)
-            tsByMessage.append(np.ones((numPoints), dtype=np.float64) * ts)
-    if not pointsByMessage: # None of the messages contained any points
+                for y in range(height):
+                    data_array[x * height + y, :] = np.frombuffer(data[ptr:ptr+dtype_fields['itemsize']], dtype=dtype_fields)
+                    ptr += PointStep
+
+            pointsByMessage_2.append(data_array)
+            tsByMessage_2.append(np.ones(numPoints, dtype=np.float64) * timestamp)  #Meter aquí la traspuesta de la matriz???
+            #Reset dtype_fields variables for the next msg
+            dtype_fields['names'] = []
+            dtype_fields['offsets'] = []
+            dtype_fields['formats'] = []
+            dtype_fields['count'] = []
+
+    if not pointsByMessage_2:  # None of the messages contained any points
         return None
-    points = np.concatenate(pointsByMessage)        
-    ts = np.concatenate(tsByMessage)        
-
+    data_array = np.concatenate(pointsByMessage_2)
+    timestamp = np.concatenate(tsByMessage_2)
     # Crop arrays to number of events
     outDict = {
-        'ts': ts,
-        'point': points,
+        'ts': timestamp,
+        'point': data_array,
         }
     return outDict