def batches()

in src/datasets.py [0:0]

• 36 lines of code
• 16 McCabe index (conditional complexity)

    def batches(self, batch_size):
        n_samples = -1

        samples_location = 0 if self.single else 1

        if DatasetKeyConstants.ray_directions_samples in self.batch_input_dirs[0]:
            n_samples = self.batch_input_dirs[0][DatasetKeyConstants.ray_directions_samples].shape[samples_location]
        elif DatasetKeyConstants.color_image_samples in self.batch_input_dirs[0]:
            n_samples = self.batch_input_dirs[0][DatasetKeyConstants.color_image_samples].shape[samples_location]
        elif DatasetKeyConstants.depth_image_samples in self.batch_input_dirs[0]:
            n_samples = self.batch_input_dirs[0][DatasetKeyConstants.depth_image_samples].shape[samples_location]

        if n_samples == -1:
            print("ERROR: unable to batch sample data!")

        for batch0 in range(0, n_samples, batch_size):
            batch_input_dirs = []
            train_targets = []

            for idx in range(len(self.batch_input_dirs)):
                inner_dir = {}

                for key in self.batch_input_dirs[idx]:
                    if key == DatasetKeyConstants.color_image_samples or \
                            key == DatasetKeyConstants.depth_image_samples or \
                            key == DatasetKeyConstants.ray_directions_samples:
                        if self.single:
                            inner_dir[key] = self.batch_input_dirs[idx][key][None, batch0:batch0 + batch_size, :]
                        else:
                            inner_dir[key] = self.batch_input_dirs[idx][key][:, batch0:batch0 + batch_size, :]
                    else:
                        if self.single:
                            inner_dir[key] = self.batch_input_dirs[idx][key][None]
                        else:
                            inner_dir[key] = self.batch_input_dirs[idx][key]

                batch_input_dirs.append(inner_dir)

            for idx in range(len(self.train_targets)):
                if self.single:
                    train_targets.append(self.train_targets[idx][batch0:batch0 + batch_size])
                else:
                    train_targets.append(self.train_targets[idx][0, batch0:batch0 + batch_size])

            yield SampleDataWrapper(batch_input_dirs, train_targets, False)