src/features.py [313:332]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    def prepare_batch(self, data, config):
        image_indices = data[DatasetKeyConstants.image_sample_indices]

        ret_dict = {DatasetKeyConstants.image_pose: data[DatasetKeyConstants.image_pose][0],
                    DatasetKeyConstants.image_rotation: data[DatasetKeyConstants.image_rotation][0],
                    DatasetKeyConstants.ray_directions_samples: data[DatasetKeyConstants.ray_directions][image_indices]}

        if DatasetKeyConstants.depth_image_full in data:
            depth_image = data[DatasetKeyConstants.depth_image_full][0]
            depth_image = depth_image.reshape(depth_image.shape[0] * depth_image.shape[1], depth_image.shape[2])
            depth_image = depth_image[image_indices]

            ret_dict[DatasetKeyConstants.depth_image_samples] = depth_image

        return ret_dict

    def batch(self, data, **kwargs) -> (torch.tensor, torch.tensor):
        poses = data[DatasetKeyConstants.image_pose]
        rotations = data[DatasetKeyConstants.image_rotation]
        directions = data[DatasetKeyConstants.ray_directions_samples]
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



src/features.py [686:705]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    def prepare_batch(self, data, config):
        image_indices = data[DatasetKeyConstants.image_sample_indices]

        ret_dict = {DatasetKeyConstants.image_pose: data[DatasetKeyConstants.image_pose][0],
                    DatasetKeyConstants.image_rotation: data[DatasetKeyConstants.image_rotation][0],
                    DatasetKeyConstants.ray_directions_samples: data[DatasetKeyConstants.ray_directions][image_indices]}

        if DatasetKeyConstants.depth_image_full in data:
            depth_image = data[DatasetKeyConstants.depth_image_full][0]
            depth_image = depth_image.reshape(depth_image.shape[0] * depth_image.shape[1], depth_image.shape[2])
            depth_image = depth_image[image_indices]

            ret_dict[DatasetKeyConstants.depth_image_samples] = depth_image

        return ret_dict

    def batch(self, data, **kwargs) -> (torch.tensor, torch.tensor):
        poses = data[DatasetKeyConstants.image_pose]
        rotations = data[DatasetKeyConstants.image_rotation]
        directions = data[DatasetKeyConstants.ray_directions_samples]
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -