def get_trajectory_Rt()

in datasets/replica.py [0:0]

• 25 lines of code
• 7 McCabe index (conditional complexity)

    def get_trajectory_Rt(self):
        Rt = []
        for idxstr in self.seq_idxs:
            episode_Rt = []
            episode_path = os.path.join(self.datapath, idxstr)
            with open(os.path.join(episode_path, 'cameras.json'), 'r') as f:
                cameras = json.load(f)

                for i in range(0, self.episode_len):
                    episode_Rt.append(torch.Tensor(cameras[i]['Rt']))
            episode_Rt = torch.stack(episode_Rt, dim=0)

            trim = episode_Rt.shape[0] % (self.seq_len * self.step)
            episode_Rt = episode_Rt[: episode_Rt.shape[0] - trim]
            Rt.append(episode_Rt)

        Rt = torch.stack(Rt, dim=0)

        # this basically samples points at the stride length
        Rt = Rt.view(-1, self.seq_len, self.step, 4, 4).permute(0, 2, 1, 3, 4).reshape(-1, self.seq_len, 4, 4)

        if self.center is not None:
            Rt = normalize_trajectory(Rt, center=self.center, normalize_rotation=self.normalize_rotation)

        if self.single_sample_per_trajectory:
            # randomly select a single point along each trajectory
            selected_indices = torch.multinomial(torch.ones(Rt.shape[:2]), num_samples=1).squeeze()
            bool_mask = torch.eye(self.seq_len)[selected_indices].bool()
            Rt = Rt[bool_mask].unsqueeze(1)

        if self.rot_aug:
            for i in range(Rt.shape[0]):
                Rt[i] = random_rotation_augment(Rt[i])
        return Rt