| 24 | |
| 25 | |
| 26 | class MVRecon(pl.LightningModule): |
| 27 | def __init__( |
| 28 | self, |
| 29 | lrm_generator_config, |
| 30 | input_size=256, |
| 31 | render_size=512, |
| 32 | init_ckpt=None, |
| 33 | ): |
| 34 | super(MVRecon, self).__init__() |
| 35 | |
| 36 | self.input_size = input_size |
| 37 | self.render_size = render_size |
| 38 | |
| 39 | # init modules |
| 40 | self.lrm_generator = instantiate_from_config(lrm_generator_config) |
| 41 | |
| 42 | self.lpips = LearnedPerceptualImagePatchSimilarity(net_type='vgg') |
| 43 | |
| 44 | # Load weights from pretrained MVRecon model, and use the mlp |
| 45 | # weights to initialize the weights of sdf and rgb mlps. |
| 46 | if init_ckpt is not None: |
| 47 | sd = torch.load(init_ckpt, map_location='cpu')['state_dict'] |
| 48 | sd = {k: v for k, v in sd.items() if k.startswith('lrm_generator')} |
| 49 | sd_fc = {} |
| 50 | for k, v in sd.items(): |
| 51 | if k.startswith('lrm_generator.synthesizer.decoder.net.'): |
| 52 | if k.startswith('lrm_generator.synthesizer.decoder.net.6.'): # last layer |
| 53 | # Here we assume the density filed's isosurface threshold is t, |
| 54 | # we reverse the sign of density filed to initialize SDF field. |
| 55 | # -(w*x + b - t) = (-w)*x + (t - b) |
| 56 | if 'weight' in k: |
| 57 | sd_fc[k.replace('net.', 'net_sdf.')] = -v[0:1] |
| 58 | else: |
| 59 | sd_fc[k.replace('net.', 'net_sdf.')] = 10.0 - v[0:1] |
| 60 | sd_fc[k.replace('net.', 'net_rgb.')] = v[1:4] |
| 61 | else: |
| 62 | sd_fc[k.replace('net.', 'net_sdf.')] = v |
| 63 | sd_fc[k.replace('net.', 'net_rgb.')] = v |
| 64 | else: |
| 65 | sd_fc[k] = v |
| 66 | sd_fc = {k.replace('lrm_generator.', ''): v for k, v in sd_fc.items()} |
| 67 | # missing `net_deformation` and `net_weight` parameters |
| 68 | self.lrm_generator.load_state_dict(sd_fc, strict=False) |
| 69 | print(f'Loaded weights from {init_ckpt}') |
| 70 | |
| 71 | self.validation_step_outputs = [] |
| 72 | |
| 73 | def on_fit_start(self): |
| 74 | device = torch.device(f'cuda:{self.global_rank}') |
| 75 | self.lrm_generator.init_flexicubes_geometry(device) |
| 76 | if self.global_rank == 0: |
| 77 | os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True) |
| 78 | os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True) |
| 79 | |
| 80 | def prepare_batch_data(self, batch): |
| 81 | lrm_generator_input = {} |
| 82 | render_gt = {} |
| 83 |
nothing calls this directly
no outgoing calls
no test coverage detected