Method forward

model/model_utils.py:277–369

(self, feats, timestamp=None)

    @autocast('cuda', enabled=False)
    def forward(self, feats, timestamp=None):
        # [B, V, N, D]
        assert feats.shape[-1] == self.embed_dim
        B = feats.shape[0]
        input_views = feats.shape[1]
        N = feats.shape[1] * feats.shape[2] * self.ratio * self.n_sample
        feats = feats.type(torch.float32)
        feats = rearrange(feats, 'b v n d -> (b v) n d') 

        gsparams = {}
        prior_params = {}
        for key in self.fix_keys:
            if key == "scale":
                fix_v = 0.03 * torch.ones(B, N, 3).to(feats.device)
            elif key == "rot_static":
                fix_v = torch.zeros(B, N, 4).to(feats.device)
                fix_v[:, :, 0] = 1.
            elif key == "rot_dynamic":
                fix_v = torch.zeros(B, N, 4).to(feats.device)
            elif key == "xyz_dynamic":
                fix_v = torch.zeros(B, N, self.key_dims[key]//3, 3).to(feats.device)
            elif key == "opacity_dynamic":
                fix_v = torch.ones(B, N, 1).to(feats.device)
            else:
                raise NotImplementedError
            gsparams[key] = fix_v
        
        def reorder(v):
            # [B*V, H*W, d*r*r]
            v = rearrange(v, 'B (h w) D -> B D h w', h=self.actual_h, w=self.actual_w).contiguous()
            v = self.pixelshuffle(v)  # [B*V, d, H*r, W*r]
            v = rearrange(v, '(b v) d h w -> b (v h w) d', v=input_views)
            return v
        
        for key in self.pred_keys:
            v = feats
            v = self.gs_layer[key](v)  
            v = reorder(v)
            
            gsparams[key] = self.key_activation(v, key)
        
        for key in self.sample_keys:
            logits_pred = torch.ones(feats.shape[0], feats.shape[1], 1).to(feats.device).float()
            activation = "tanh"
            means = reorder(self.mix_layer[f"{key}_mean"](feats))  # [B, N, nr_mix * dim]
            log_scales = reorder(self.mix_layer[f"{key}_scale"](feats))
            logits, means, log_scales = self.prior.expand_params(logits_pred, means, log_scales, mean_activation=activation)
            prior_params[key] = {"logits": logits, "means": means, "log_scales": log_scales}  # [B, N*r, dim, nr_mix]
            
            if key == "xyz_static":
                val, probs = self.prior.sample(logits, means, log_scales)
                x_pred, y_pred, z_pred = val.chunk(3, dim=-1)
                y_val = (0.5 + y_pred * 0.5)
                assert y_val.min() >= 0. and y_val.max() <= 1., f"y_val: {y_val.min()}, {y_val.max()}"
                x_offset = self.x_max * x_pred
                z_offset = self.z_max * z_pred
                x_map = self.x_map.repeat(self.opt.input_frames).reshape(1, -1, 1)
                z_map = self.z_map.repeat(self.opt.input_frames).reshape(1, -1, 1)