Class for single ScanNet scene processing. Single scene file contains multiple RGB-D images.
| 50 | |
| 51 | |
| 52 | class SensorData: |
| 53 | """Class for single ScanNet scene processing. |
| 54 | |
| 55 | Single scene file contains multiple RGB-D images. |
| 56 | """ |
| 57 | |
| 58 | def __init__(self, filename, fast=False): |
| 59 | self.version = 4 |
| 60 | self.load(filename, fast) |
| 61 | |
| 62 | def load(self, filename, fast): |
| 63 | """Load a single scene data with multiple RGBD frames.""" |
| 64 | with open(filename, 'rb') as f: |
| 65 | version = struct.unpack('I', f.read(4))[0] |
| 66 | assert self.version == version |
| 67 | strlen = struct.unpack('Q', f.read(8))[0] |
| 68 | self.sensor_name = b''.join( |
| 69 | struct.unpack('c' * strlen, f.read(strlen))) |
| 70 | self.intrinsic_color = np.asarray(struct.unpack( |
| 71 | 'f' * 16, f.read(16 * 4)), |
| 72 | dtype=np.float32).reshape(4, 4) |
| 73 | self.extrinsic_color = np.asarray(struct.unpack( |
| 74 | 'f' * 16, f.read(16 * 4)), |
| 75 | dtype=np.float32).reshape(4, 4) |
| 76 | self.intrinsic_depth = np.asarray(struct.unpack( |
| 77 | 'f' * 16, f.read(16 * 4)), |
| 78 | dtype=np.float32).reshape(4, 4) |
| 79 | self.extrinsic_depth = np.asarray(struct.unpack( |
| 80 | 'f' * 16, f.read(16 * 4)), |
| 81 | dtype=np.float32).reshape(4, 4) |
| 82 | self.color_compression_type = COMPRESSION_TYPE_COLOR[struct.unpack( |
| 83 | 'i', f.read(4))[0]] |
| 84 | self.depth_compression_type = COMPRESSION_TYPE_DEPTH[struct.unpack( |
| 85 | 'i', f.read(4))[0]] |
| 86 | self.color_width = struct.unpack('I', f.read(4))[0] |
| 87 | self.color_height = struct.unpack('I', f.read(4))[0] |
| 88 | self.depth_width = struct.unpack('I', f.read(4))[0] |
| 89 | self.depth_height = struct.unpack('I', f.read(4))[0] |
| 90 | self.depth_shift = struct.unpack('f', f.read(4))[0] |
| 91 | num_frames = struct.unpack('Q', f.read(8))[0] |
| 92 | self.num_frames = num_frames |
| 93 | self.frames = [] |
| 94 | if fast: |
| 95 | index = list(range(num_frames))[::10] |
| 96 | else: |
| 97 | index = list(range(num_frames)) |
| 98 | self.index = index |
| 99 | for i in range(num_frames): |
| 100 | frame = RGBDFrame() |
| 101 | frame.load(f) |
| 102 | if i in index: |
| 103 | self.frames.append(frame) |
| 104 | |
| 105 | def export_depth_images(self, output_path): |
| 106 | """Export depth images to the output path.""" |
| 107 | if not os.path.exists(output_path): |
| 108 | os.makedirs(output_path) |
| 109 | for f in range(len(self.frames)): |