| 3 | |
| 4 | |
| 5 | class Criterion(nn.Module): |
| 6 | def __init__(self, args) -> None: |
| 7 | super().__init__() |
| 8 | self.args = args |
| 9 | self.rgb_weight = args.criteria["rgb_weight"] |
| 10 | self.depth_weight = args.criteria["depth_weight"] |
| 11 | self.sdf_weight = args.criteria["sdf_weight"] |
| 12 | self.fs_weight = args.criteria["fs_weight"] |
| 13 | self.truncation = args.criteria["sdf_truncation"] |
| 14 | self.max_dpeth = args.data_specs["max_depth"] |
| 15 | |
| 16 | def forward(self, outputs, obs, use_color_loss=True, |
| 17 | use_depth_loss=True, compute_sdf_loss=True, |
| 18 | weight_depth_loss=False): |
| 19 | |
| 20 | img, depth = obs # img:(1,N_rays,3) depth:(1,N_rays) |
| 21 | loss = 0 |
| 22 | loss_dict = {} |
| 23 | |
| 24 | pred_depth = outputs["depth"] # (N_rays) |
| 25 | pred_color = outputs["color"] # (N_rays,3) |
| 26 | pred_sdf = outputs["sdf"] # (N_rays,N_samples_every_ray) |
| 27 | z_vals = outputs["z_vals"] # (N_rays,N_samples_every_ray) |
| 28 | ray_mask = outputs["ray_mask"] # (1,N_rays) |
| 29 | weights = outputs["weights"] # (N_rays,N_samples_every_ray) |
| 30 | |
| 31 | gt_depth = depth[ray_mask] # (N_rays) |
| 32 | gt_color = img[ray_mask] # (N_rays,3) |
| 33 | |
| 34 | if use_depth_loss: |
| 35 | valid_depth = (gt_depth > 0.01) & (gt_depth < self.max_dpeth) |
| 36 | depth_loss = (gt_depth - pred_depth).abs() # (N_rays) |
| 37 | |
| 38 | if weight_depth_loss: |
| 39 | depth_var = weights * ((pred_depth.unsqueeze(-1) - z_vals) ** 2) |
| 40 | depth_var = torch.sum(depth_var, -1) |
| 41 | tmp = depth_loss / torch.sqrt(depth_var + 1e-10) |
| 42 | valid_depth = (tmp < 10 * tmp.median()) & valid_depth |
| 43 | depth_loss = depth_loss[valid_depth].mean() |
| 44 | loss += self.depth_weight * depth_loss |
| 45 | # loss_dict["depth_loss"] = self.depth_weight * depth_loss.item() |
| 46 | |
| 47 | if use_color_loss: |
| 48 | color_loss = (gt_color - pred_color).abs()[valid_depth].mean() |
| 49 | loss += self.rgb_weight * color_loss |
| 50 | # loss_dict["color_loss"] = self.rgb_weight * color_loss.item() |
| 51 | |
| 52 | if compute_sdf_loss: |
| 53 | fs_loss, sdf_loss = self.get_sdf_loss( |
| 54 | z_vals, gt_depth, pred_sdf, |
| 55 | truncation=self.truncation, |
| 56 | loss_type='l2' |
| 57 | ) |
| 58 | loss += self.fs_weight * fs_loss |
| 59 | loss += self.sdf_weight * sdf_loss |
| 60 | # loss_dict["fs_loss"] = self.fs_weight * fs_loss.item() |
| 61 | # loss_dict["bs_loss"] = back_loss.item() |
| 62 | # loss_dict["sdf_loss"] = self.sdf_weight * sdf_loss.item() |