Entry for running compression.
(self, step: int)
| 1257 | |
| 1258 | @torch.no_grad() |
| 1259 | def run_compression(self, step: int): |
| 1260 | """Entry for running compression.""" |
| 1261 | print("Running compression...") |
| 1262 | world_rank = self.world_rank |
| 1263 | |
| 1264 | compress_dir = f"{cfg.result_dir}/compression/rank{world_rank}" |
| 1265 | |
| 1266 | if os.path.exists(compress_dir): |
| 1267 | shutil.rmtree(compress_dir) |
| 1268 | os.makedirs(compress_dir) |
| 1269 | |
| 1270 | self.run_param_distribution_vis(self.splats, save_dir=f"{cfg.result_dir}/visualization/raw") |
| 1271 | |
| 1272 | if isinstance(self.compression_method, PngCompression): |
| 1273 | self.compression_method.compress(compress_dir, self.splats) |
| 1274 | elif isinstance(self.compression_method, EntropyCodingCompression): |
| 1275 | self.compression_method.compress(compress_dir, self.splats, self.entropy_models) |
| 1276 | elif isinstance(self.compression_method, HevcCompression): |
| 1277 | self.compression_method.compress(compress_dir, self.splats) |
| 1278 | else: |
| 1279 | raise NotImplementedError(f"The compression method is not implemented yet.") |
| 1280 | |
| 1281 | # evaluate compression |
| 1282 | splats_c = self.compression_method.decompress(compress_dir) |
| 1283 | |
| 1284 | self.run_param_distribution_vis(splats_c, save_dir=f"{cfg.result_dir}/visualization/quant") |
| 1285 | |
| 1286 | for k in splats_c.keys(): |
| 1287 | self.splats[k].data = splats_c[k].to(self.device) |
| 1288 | self.eval(step=step, stage="compress") |
| 1289 | self.render_traj(step=step, stage="compress") |
| 1290 | |
| 1291 | @torch.no_grad() |
| 1292 | def run_post_training_compression(self, step: int): |
nothing calls this directly
no test coverage detected