Method forward

model/model_utils.py:623–673

(self, frames, depths, cond_times=None)

        return pred_gs

    def forward(self, frames, depths, cond_times=None):
        frames = torch.cat([frames[:, 0:1], frames[:, -1:]], dim=1)  # [B, 2, C, H, W]
        depths = torch.cat([depths[:, 0:1], depths[:, -1:]], dim=1)  # [B, 2, C, H, W]
        input_views = frames.shape[1]
        frames = rearrange(frames, 'b v c h w -> (b v) 1 c h w')
        depths = rearrange(depths, 'b v c h w -> (b v) 1 c h w')
        encoder_output = self.forward_encoder(frames, depths)  # [B*V, 1, N, D]
        encoder_output = rearrange(encoder_output, '(b v) 1 n d -> b v n d', v=input_views)
        encoder_0 = encoder_output[:, 0:1]  # [B, V, N, D]
        encoder_1 = encoder_output[:, -1:]
        upsampled_encoder_0 = self.gs_predictor.upsampling(encoder_0)
        upsampled_encoder_1 = self.gs_predictor.upsampling(encoder_1)
        gs_0, _ = self.gs_predictor.predictor(upsampled_encoder_0)
        gs_1, _ = self.gs_predictor.predictor(upsampled_encoder_1)

        if self.opt.use_dino:
            frames = rearrange(frames, '(b v) 1 c h w -> b v c h w', v=input_views)
            depths = rearrange(depths, '(b v) 1 c h w -> b v c h w', v=input_views)
            condition_output = self.forward_condition(frames, depths)  # [B, V, N, D]
            condition_0 = condition_output[:, 0:1]  # [B, V, N, D]
            condition_1 = condition_output[:, -1:]
        else:
            condition_0 = encoder_0
            condition_1 = encoder_1

        decoder_input = torch.cat([encoder_0, encoder_1], dim=2)  # [B, V, 2*N, D]
        condition_input = torch.cat([condition_1, condition_0], dim=2)  # [B, V, 2*N, D]
        decoder_output = self.decoder(decoder_input, condition_input)
        decoder_output_0 = decoder_output[:, :, :self.decoder.token_len]  # [B, V, N, D]
        decoder_output_1 = decoder_output[:, :, self.decoder.token_len:]

        upsampled_dynamic_0 = self.upsampling(decoder_output_0)
        upsampled_dynamic_1 = self.upsampling(decoder_output_1)
        upsampled_0 = torch.cat([self.encoder_proj(upsampled_encoder_0), upsampled_dynamic_0], dim=-1)
        upsampled_1 = torch.cat([self.encoder_proj(upsampled_encoder_1), upsampled_dynamic_1], dim=-1)
        gs_0_dynamic, gs_0_prior = self.gs_dynamic_predictor(upsampled_0)
        gs_1_dynamic, gs_1_prior = self.gs_dynamic_predictor(upsampled_1)

        for key in ["xyz_dynamic", "opacity_dynamic"]:
            gs_0[key] = gs_0_dynamic[key]
            gs_1[key] = gs_1_dynamic[key]

        gs_0 = self.combine(gs_0)
        gs_1 = self.combine(gs_1)

        # concate gs_0 and gs_1
        pred_gs = {}
        for key in gs_0.keys():
            pred_gs[key] = torch.cat([gs_0[key], gs_1[key]], dim=1)

        return {'pred_gs': pred_gs, 'gs_0': gs_0, 'gs_1': gs_1}