| 769 | return x |
| 770 | |
| 771 | class MVSNet(nn.Module): |
| 772 | def __init__(self, |
| 773 | num_groups=1, |
| 774 | norm_act=InPlaceABN, |
| 775 | levels=1): |
| 776 | super(MVSNet, self).__init__() |
| 777 | self.levels = levels # 3 depth levels |
| 778 | self.n_depths = [128,32,8] |
| 779 | self.G = num_groups # number of groups in groupwise correlation |
| 780 | self.feature = FeatureNet() |
| 781 | |
| 782 | self.N_importance = 0 |
| 783 | self.chunk = 1024 |
| 784 | |
| 785 | self.cost_reg_2 = CostRegNet(32+9, norm_act) |
| 786 | |
| 787 | def build_volume_costvar(self, feats, proj_mats, depth_values, pad=0): |
| 788 | # feats: (B, V, C, H, W) |
| 789 | # proj_mats: (B, V, 3, 4) |
| 790 | # depth_values: (B, D, H, W) |
| 791 | # cost_reg: nn.Module of input (B, C, D, h, w) and output (B, 1, D, h, w) |
| 792 | # volume_sum [B, G, D, h, w] |
| 793 | # prob_volume [B D H W] |
| 794 | # volume_feature [B C D H W] |
| 795 | |
| 796 | B, V, C, H, W = feats.shape |
| 797 | D = depth_values.shape[1] |
| 798 | |
| 799 | ref_feats, src_feats = feats[:, 0], feats[:, 1:] |
| 800 | src_feats = src_feats.permute(1, 0, 2, 3, 4) # (V-1, B, C, h, w) |
| 801 | proj_mats = proj_mats[:, 1:] |
| 802 | proj_mats = proj_mats.permute(1, 0, 2, 3) # (V-1, B, 3, 4) |
| 803 | |
| 804 | if pad > 0: |
| 805 | ref_feats = F.pad(ref_feats, (pad, pad, pad, pad), "constant", 0) |
| 806 | |
| 807 | ref_volume = ref_feats.unsqueeze(2).repeat(1, 1, D, 1, 1) # (B, C, D, h, w) |
| 808 | |
| 809 | volume_sum = ref_volume |
| 810 | volume_sq_sum = ref_volume ** 2 |
| 811 | |
| 812 | del ref_feats |
| 813 | |
| 814 | in_masks = torch.ones((B, 1, D, H + pad * 2, W + pad * 2), device=volume_sum.device) |
| 815 | for i, (src_feat, proj_mat) in enumerate(zip(src_feats, proj_mats)): |
| 816 | warped_volume, grid = homo_warp(src_feat, proj_mat, depth_values, pad=pad) |
| 817 | |
| 818 | grid = grid.view(B, 1, D, H + pad * 2, W + pad * 2, 2) |
| 819 | in_mask = ((grid > -1.0) * (grid < 1.0)) |
| 820 | in_mask = (in_mask[..., 0] * in_mask[..., 1]) |
| 821 | in_masks += in_mask.float() |
| 822 | |
| 823 | if self.training: |
| 824 | volume_sum = volume_sum + warped_volume |
| 825 | volume_sq_sum = volume_sq_sum + warped_volume ** 2 |
| 826 | else: |
| 827 | volume_sum += warped_volume |
| 828 | volume_sq_sum += warped_volume.pow_(2) |