Method forward

train.py:177–229

(self,
                ts_w,
                grid,
                encode_w,
                step=0,
                flows=None)

        return deformed_grid

    def forward(self,
                ts_w,
                grid,
                encode_w,
                step=0,
                flows=None):
        # grid -> positional encoding
        # ts_w -> embedding
        grid = rearrange(grid, 'b n c -> (b n) c')
        results_list = []
        flow_loss_list = []
        deform_list = []
        for i in range(self.num_models):
            deformed_grid = self.deform_pts(ts_w, grid, encode_w, step, i)  # [batch * num_pixels, 2]
            deform_list.append(deformed_grid)
            # Compute optical flow loss.
            flow_loss = 0
            if self.hparams.flow_loss > 0 and not self.hparams.test:
                if flows.max() > -1e2 and step > self.hparams.flow_step:
                    grid_new = grid + flows.squeeze(0)
                    deformed_grid_new = self.deform_pts(
                        ts_w + 1, grid_new, encode_w, step, i)
                    flow_loss = (deformed_grid_new, deformed_grid)
            flow_loss_list.append(flow_loss)
            if self.hparams.vid_hash:
                pe_deformed_grid = (deformed_grid + 0.3) / 1.6
            else:
                pe_deformed_grid = self.embeddings['xyz'](deformed_grid)
            if not self.training and self.hparams.canonical_dir is not None:
                w, h = self.img_wh
                canonical_img = self.canonical_img.squeeze(0)
                h_c, w_c = canonical_img.shape[1:3]
                grid_new = deformed_grid.clone()
                grid_new[..., 1] = (2 * deformed_grid[..., 0] - 1) * h / h_c
                grid_new[..., 0] = (2 * deformed_grid[..., 1] - 1) * w / w_c
                if len(canonical_img.shape) == 3:
                    canonical_img = canonical_img.unsqueeze(0)
                results = torch.nn.functional.grid_sample(
                    canonical_img[i:i + 1].permute(0, 3, 1, 2),
                    grid_new.unsqueeze(1).unsqueeze(0),
                    mode='bilinear',
                    padding_mode='border')
                results = results.squeeze().permute(1,0)
            else:
                results = self.models[f'implicit_video_{i}'](pe_deformed_grid)

            results_list.append(results)

        ret = edict(rgbs=results_list,
                    flow_loss=flow_loss_list,
                    deform=deform_list)

        return ret

    def setup(self, stage):
        if not self.hparams.test: