def map_dataset_sample()

in habitat_baselines/il/data/nav_data.py [0:0]

• 92 lines of code
• 14 McCabe index (conditional complexity)

    def map_dataset_sample(self, x: Dict) -> Tuple:
        """Mapper function to pre-process webdataset sample, example:
        img features, planner & controller actions etc.
        Args:
            x: webdataset sample containing ep_id, question, answer and imgs
        Returns:
            Processed sample containing img features, planner & controller actions etc.
        """
        idx = x["episode_id"]
        question = x["question"]
        answer = x["answer"]

        if len(question) < self.max_q_len:
            diff = self.max_q_len - len(question)
            for _ in range(diff):
                question.append(0)

        self.frame_queue = np.array(
            [img.transpose(2, 0, 1) / 255.0 for img in list(x.values())[4:]]
        )
        self.frame_queue = torch.Tensor(self.frame_queue).to(self.device)

        if self.mode == "val":
            # works only with batch size 1
            actions = self.episodes[idx].actions
            action_length = self.episodes[idx].action_length
            scene = self.episodes[idx].scene_id
            if scene != self.config.SIMULATOR.SCENE:
                logger.info("[ Loading scene - {}]".format(scene))
                self.config.defrost()
                self.config.SIMULATOR.SCENE = scene
                self.config.freeze()
                self.env.sim.reconfigure(self.config.SIMULATOR)

            goal_pos = self.episodes[idx].goals[0].position

            return idx, question, answer, actions, action_length, goal_pos

        planner_actions = self.episodes[idx].planner_actions
        controller_actions = self.episodes[idx].controller_actions

        planner_hidden_idx = self.episodes[idx].planner_hidden_idx

        planner_action_length = self.episodes[idx].planner_action_length
        controller_action_length = self.episodes[idx].controller_action_length

        raw_img_feats = (
            self.get_img_features(self.frame_queue).cpu().numpy().copy()
        )
        img_feats = np.zeros(
            (self.max_action_len, raw_img_feats.shape[1]), dtype=np.float32
        )
        img_feats[: raw_img_feats.shape[0], :] = raw_img_feats.copy()

        planner_pos_queue_idx = self.episodes[idx].planner_pos_queue_idx
        controller_pos_queue_idx = self.episodes[idx].controller_pos_queue_idx

        planner_img_feats = np.zeros(
            (self.max_action_len, img_feats.shape[1]), dtype=np.float32
        )

        planner_img_feats[
            : self.episodes[idx].planner_action_length
        ] = img_feats[tuple(planner_pos_queue_idx)]

        planner_actions_in = planner_actions.clone() - 1
        planner_actions_out = planner_actions[1:].clone() - 2

        planner_actions_in[planner_action_length:].fill_(0)
        planner_mask = planner_actions_out.clone().gt(-1)

        if len(planner_actions_out) > planner_action_length:
            planner_actions_out[planner_action_length:].fill_(0)

        controller_img_feats = np.zeros(
            (self.max_action_len, img_feats.shape[1]), dtype=np.float32
        )
        controller_img_feats[:controller_action_length] = img_feats[
            tuple(controller_pos_queue_idx)
        ]

        controller_actions_in = self.episodes[idx].actions.clone() - 2

        if len(controller_actions_in) > controller_action_length:
            controller_actions_in[controller_action_length:].fill_(0)

        controller_out = controller_actions
        controller_mask = controller_out.clone().gt(-1)
        if len(controller_out) > controller_action_length:
            controller_out[controller_action_length:].fill_(0)

        # zero out forced controller return
        for i in range(controller_action_length):
            if (
                i >= self.max_controller_actions - 1
                and controller_out[i] == 0
                and (
                    self.max_controller_actions == 1
                    or controller_out[
                        i - self.max_controller_actions + 1 : i
                    ].sum()
                    == self.max_controller_actions - 1
                )
            ):
                controller_mask[i] = 0
        return (
            idx,
            question,
            answer,
            planner_img_feats,
            planner_actions_in,
            planner_actions_out,
            planner_action_length,
            planner_mask,
            controller_img_feats,
            controller_actions_in,
            planner_hidden_idx,
            controller_out,
            controller_action_length,
            controller_mask,
        )