(self, data, step_ratio=0.0)
| 66 | return decoder_out |
| 67 | |
| 68 | def forward(self, data, step_ratio=0.0): |
| 69 | # data: [B, 2, C, H, W] |
| 70 | input_frames = data['frames'][:, 0:1] # [B, 1, C, H, W], input features |
| 71 | input_depths = data['depths'][:, 0:1] # [B, 1, C, H, W], input features |
| 72 | |
| 73 | results = {} |
| 74 | |
| 75 | # predict gaussians |
| 76 | decoder_out = self.forward_gaussians(input_frames, input_depths) |
| 77 | with autocast('cuda', enabled=False): |
| 78 | render_pkg = self.gaussian_renderer(decoder_out["pred_gs"], self.background, opt=self.opt) |
| 79 | output_frames = render_pkg["render"] |
| 80 | pred_depths = render_pkg["depth"] |
| 81 | mse_loss = F.mse_loss(output_frames, input_frames) |
| 82 | loss = mse_loss |
| 83 | |
| 84 | if self.opt.depth_downsample: |
| 85 | actual_h = int(self.opt.down_resolution[0] * (2 ** self.opt.decoder_ratio / self.opt.patch_size)) |
| 86 | actual_w = int(self.opt.down_resolution[1] * (2 ** self.opt.decoder_ratio / self.opt.patch_size)) |
| 87 | views = input_depths.shape[1] |
| 88 | pred_depths = rearrange(pred_depths, 'b v c h w -> (b v) c h w') |
| 89 | input_depths = rearrange(input_depths, 'b v c h w -> (b v) c h w') |
| 90 | pred_depths = F.interpolate(pred_depths, (actual_h, actual_w), mode='bilinear', align_corners=True) |
| 91 | input_depths = F.interpolate(input_depths, (actual_h, actual_w), mode='nearest') |
| 92 | pred_depths = rearrange(pred_depths, '(b v) c h w -> b v c h w', v=views) |
| 93 | input_depths = rearrange(input_depths, '(b v) c h w -> b v c h w', v=views) |
| 94 | |
| 95 | depth_loss = torch.zeros(1, device=input_frames.device) |
| 96 | if self.opt.epoch > self.opt.depth_start_epoch: |
| 97 | loss_func = ssitrim_loss if "trim" in self.opt.depth_loss_type else ssimse_loss |
| 98 | depth_loss = loss_func(pred_depths, input_depths, None, self.opt.ignore_large_loss) # normalized input_depths with no mask |
| 99 | |
| 100 | loss = loss + self.opt.lambda_depth * depth_loss |
| 101 | |
| 102 | if self.opt.lambda_lpips > 0 and self.opt.epoch > self.opt.lpips_start_epoch: |
| 103 | down_res_H, down_res_W = self.opt.down_resolution |
| 104 | loss_lpips = self.lpips_loss( |
| 105 | F.interpolate(input_frames.reshape(-1, 3, down_res_H, down_res_W) * 2 - 1, (256, 256), mode='bilinear', align_corners=False), |
| 106 | F.interpolate(output_frames.reshape(-1, 3, down_res_H, down_res_W) * 2 - 1, (256, 256), mode='bilinear', align_corners=False), |
| 107 | ).mean() |
| 108 | results['loss_lpips'] = loss_lpips |
| 109 | loss = loss + self.opt.lambda_lpips * loss_lpips |
| 110 | |
| 111 | pred_depth = 1.0 / (render_pkg["depth"] + 1e-8) |
| 112 | B, V, C, H, W = pred_depth.shape |
| 113 | reshaped_depth = pred_depth.view(B, V * C * H * W) # Shape [B*V, H*W] |
| 114 | min_vals = reshaped_depth.min(dim=1, keepdim=True)[0] # Shape [B*V, 1] |
| 115 | max_vals = reshaped_depth.max(dim=1, keepdim=True)[0] # Shape [B*V, 1] |
| 116 | # Normalize the depth values |
| 117 | pred_depth = (pred_depth - min_vals.view(B, 1, 1, 1, 1)) / (max_vals.view(B, 1, 1, 1, 1) - min_vals.view(B, 1, 1, 1, 1) + 1e-8) |
| 118 | pred_depth = pred_depth.clamp(0, 1) # Ensure values are between 0 and 1 |
| 119 | |
| 120 | results['loss'] = loss |
| 121 | results['mse_loss'] = mse_loss |
| 122 | results['depth_loss'] = depth_loss |
| 123 | results['pred_frames'] = output_frames |
| 124 | results['gaussians'] = decoder_out["pred_gs"] |
| 125 | results['pred_depths'] = pred_depth |
nothing calls this directly
no test coverage detected