| 2165 | return rgbd_image |
| 2166 | |
| 2167 | def chunk(self, chunks: int, dim: int = 0) -> T.List['RGBDImage']: |
| 2168 | out_dict = dict() |
| 2169 | total = None |
| 2170 | for attr_name in ['rgb', 'depth', 'normal_w', 'hit_map', 'camera', 'feature']: |
| 2171 | arr = getattr(self, attr_name, None) |
| 2172 | if arr is not None: |
| 2173 | chunked_arr = arr.chunk(chunks=chunks, dim=dim) |
| 2174 | out_dict[attr_name] = chunked_arr |
| 2175 | if total is None: |
| 2176 | total = len(chunked_arr) |
| 2177 | else: |
| 2178 | assert len(chunked_arr) == total |
| 2179 | rgbd_images = [] |
| 2180 | for i in range(total): |
| 2181 | d = dict() |
| 2182 | for attr_name in out_dict: |
| 2183 | d[attr_name] = out_dict[attr_name][i] |
| 2184 | rgbd_image = RGBDImage(**d) |
| 2185 | rgbd_images.append(rgbd_image) |
| 2186 | return rgbd_images |
| 2187 | |
| 2188 | def to(self, device: torch.device) -> 'RGBDImage': |
| 2189 | for attr_name in ['rgb', 'depth', 'normal_w', 'hit_map', 'camera', 'feature']: |