def __getitem__()

in datasets/vizdoom.py [0:0]

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

    def __getitem__(self, idx):
        random.seed()

        if self.samples_per_epoch:
            idx = random.randint(0, len(self.seq_idxs) - 1)
            idxstr = self.seq_idxs[idx]
            # Choose random start point (first ~50 frames have no movement)
            seq_start = random.randint(50, self.episode_len + 50 - (self.seq_len * self.step))
        else:
            seq_idx = math.floor(idx / (self.episode_len / (self.seq_len * self.step)))
            seq_idx = int(self.seq_idxs[seq_idx])
            seq_start = ((idx % (self.episode_len / (self.seq_len * self.step))) * (self.seq_len * self.step)) + 50
            seq_start = int(seq_start)
            idxstr = str(seq_idx).zfill(2)

        # Load cameras
        episode_path = os.path.join(self.datapath, idxstr)
        with open(os.path.join(episode_path, 'cameras.json'), 'r') as f:
            cameras = json.load(f)

        Rt = []
        K = []
        rgb = []
        depth = []

        for idx, i in enumerate(range(seq_start, seq_start + (self.seq_len * self.step), self.step)):
            Rt.append(torch.Tensor(cameras[i]['Rt']))
            K.append(torch.Tensor(cameras[i]['K']))

            _rgb = os.path.join(episode_path, str(i).zfill(3) + '_rgb.png')
            _rgb = self.resize_transform_rgb(Image.open(_rgb))
            rgb.append(_rgb)

            if self.depth:
                _depth = os.path.join(episode_path, str(i).zfill(3) + '_depth.png')
                # We dont want to normalize depth values
                _depth = self.resize_transform_depth(Image.open(_depth))
                depth.append(torch.from_numpy(np.array(_depth)).unsqueeze(0))

        rgb = torch.stack(rgb)
        depth = torch.stack(depth).float()
        K = torch.stack(K)
        Rt = torch.stack(Rt)

        Rt = Rt.unsqueeze(0)  # add batch dimension
        Rt = normalize_trajectory(Rt, center=self.center, normalize_rotation=self.normalize_rotation)
        Rt = Rt[0]  # remove batch dimension

        if self.single_sample_per_trajectory:
            selected_indices = torch.multinomial(torch.ones(Rt.shape[0]), num_samples=1).squeeze()
            rgb = rgb[selected_indices].unsqueeze(0)
            depth = depth[selected_indices].unsqueeze(0)
            K = K[selected_indices].unsqueeze(0)
            Rt = Rt[selected_indices].unsqueeze(0)

        if self.rot_aug:
            Rt = random_rotation_augment(Rt)

        # Normalize K to img_res
        downsampling_ratio = self.img_res / 480
        K[:, 0, 0] = K[:, 0, 0] * downsampling_ratio
        K[:, 1, 1] = K[:, 1, 1] * downsampling_ratio
        depth = depth / 14 * 100  # recommended scaling from game engine units to real world units

        if self.depth:
            sample = {'rgb': rgb, 'depth': depth, 'K': K, 'Rt': Rt}
        else:
            sample = {'rgb': rgb, 'K': K, 'Rt': Rt}

        return sample