| 43 | |
| 44 | |
| 45 | class SensorData: |
| 46 | |
| 47 | def __init__(self, filename): |
| 48 | self.version = 4 |
| 49 | self.load(filename) |
| 50 | |
| 51 | |
| 52 | def load(self, filename): |
| 53 | with open(filename, 'rb') as f: |
| 54 | version = struct.unpack('I', f.read(4))[0] |
| 55 | assert self.version == version |
| 56 | strlen = struct.unpack('Q', f.read(8))[0] |
| 57 | self.sensor_name = b''.join(struct.unpack('c'*strlen, f.read(strlen))) |
| 58 | self.intrinsic_color = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) |
| 59 | self.extrinsic_color = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) |
| 60 | self.intrinsic_depth = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) |
| 61 | self.extrinsic_depth = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) |
| 62 | self.color_compression_type = COMPRESSION_TYPE_COLOR[struct.unpack('i', f.read(4))[0]] |
| 63 | self.depth_compression_type = COMPRESSION_TYPE_DEPTH[struct.unpack('i', f.read(4))[0]] |
| 64 | self.color_width = struct.unpack('I', f.read(4))[0] |
| 65 | self.color_height = struct.unpack('I', f.read(4))[0] |
| 66 | self.depth_width = struct.unpack('I', f.read(4))[0] |
| 67 | self.depth_height = struct.unpack('I', f.read(4))[0] |
| 68 | self.depth_shift = struct.unpack('f', f.read(4))[0] |
| 69 | num_frames = struct.unpack('Q', f.read(8))[0] |
| 70 | self.frames = [] |
| 71 | for i in range(num_frames): |
| 72 | frame = RGBDFrame() |
| 73 | frame.load(f) |
| 74 | self.frames.append(frame) |
| 75 | |
| 76 | |
| 77 | def export_depth_images(self, output_path, image_size=None, frame_skip=1): |
| 78 | if not os.path.exists(output_path): |
| 79 | os.makedirs(output_path) |
| 80 | # print 'exporting', len(self.frames)//frame_skip, ' depth frames to', output_path |
| 81 | for f in range(0, len(self.frames), frame_skip): |
| 82 | depth_data = self.frames[f].decompress_depth(self.depth_compression_type) |
| 83 | depth = np.fromstring(depth_data, dtype=np.uint16).reshape(self.depth_height, self.depth_width) |
| 84 | if image_size is not None: |
| 85 | depth = cv2.resize(depth, (image_size[1], image_size[0]), interpolation=cv2.INTER_NEAREST) |
| 86 | imageio.imwrite(os.path.join(output_path, str(f) + '.png'), depth) |
| 87 | |
| 88 | |
| 89 | def export_color_images(self, output_path, image_size=None, frame_skip=1): |
| 90 | if not os.path.exists(output_path): |
| 91 | os.makedirs(output_path) |
| 92 | # print 'exporting', len(self.frames)//frame_skip, 'color frames to', output_path |
| 93 | for f in range(0, len(self.frames), frame_skip): |
| 94 | color = self.frames[f].decompress_color(self.color_compression_type) |
| 95 | if image_size is not None: |
| 96 | color = cv2.resize(color, (image_size[1], image_size[0]), interpolation=cv2.INTER_NEAREST) |
| 97 | imageio.imwrite(os.path.join(output_path, str(f) + '.jpg'), color) |
| 98 | |
| 99 | |
| 100 | def save_mat_to_file(self, matrix, filename): |
| 101 | with open(filename, 'w') as f: |
| 102 | for line in matrix: |