How to use the csbdeep.predict.predict_direct function in csbdeep

x = _permute_axes(img)
        channel = axes_dict(self.config.axes)['C']

        self.config.n_channel_in == x.shape[channel] or _raise(ValueError())

        # normalize
        x = normalizer.before(x,self.config.axes)
        # resize: make divisible by power of 2 to allow downsampling steps in unet
        div_n = 2 ** self.config.unet_n_depth
        x = resizer.before(x,div_n,exclude=channel)

        done = False
        while not done:
            try:
                if n_tiles == 1:
                    x = predict_direct(self.keras_model,x,channel_in=channel,channel_out=channel)
                else:
                    overlap = tile_overlap(self.config.unet_n_depth, self.config.unet_kern_size)
                    x = predict_tiled(self.keras_model,x,channel_in=channel,channel_out=channel,
                                      n_tiles=n_tiles,block_size=div_n,tile_overlap=overlap)
                done = True
            except tf.errors.ResourceExhaustedError:
                n_tiles = max(4, 2*n_tiles)
                print('Out of memory, retrying with n_tiles = %d' % n_tiles)

        n_channel_predicted = self.config.n_channel_out * (2 if self.config.probabilistic else 1)
        x.shape[channel] == n_channel_predicted or _raise(ValueError())

        x = resizer.after(x,exclude=channel)

        mean, scale = self._mean_and_scale_from_prediction(x,axis=channel)

div_n = 2 ** self.config.unet_n_depth
        x_scaled = resizer.before(x_scaled,div_n,exclude=channel)

        # move channel to the end
        x_scaled = np.moveaxis(x_scaled, channel, -1)
        channel = -1

        # u1: first rotation and prediction
        x_rot1   = self._rotate(x_scaled, axis=1, copy=False)
        u_rot1   = predict_direct(self.keras_model, x_rot1, channel_in=channel, channel_out=channel, single_sample=False,
                                  batch_size=batch_size, verbose=0)
        u1       = self._rotate(u_rot1, -1, axis=1, copy=False)

        # u2: second rotation and prediction
        x_rot2   = self._rotate(self._rotate(x_scaled, axis=2, copy=False), axis=0, copy=False)
        u_rot2   = predict_direct(self.keras_model, x_rot2, channel_in=channel, channel_out=channel, single_sample=False,
                                  batch_size=batch_size, verbose=0)
        u2       = self._rotate(self._rotate(u_rot2, -1, axis=0, copy=False), -1, axis=2, copy=False)

        n_channel_predicted = self.config.n_channel_out * (2 if self.config.probabilistic else 1)
        u_rot1.shape[channel] == n_channel_predicted or _raise(ValueError())
        u_rot2.shape[channel] == n_channel_predicted or _raise(ValueError())

        # move channel back to the front
        u1 = np.moveaxis(u1, channel, 0)
        u2 = np.moveaxis(u2, channel, 0)
        channel = 0

        # resize after prediction
        u1 = resizer.after(u1,exclude=channel)
        u2 = resizer.after(u2,exclude=channel)

x = normalizer.before(x,axes_tmp)

        # scale z up (second axis)
        x_scaled = scale_z(x,factor)

        # resize: make (x,y,z) image dimensions divisible by power of 2 to allow downsampling steps in unet
        div_n = 2 ** self.config.unet_n_depth
        x_scaled = resizer.before(x_scaled,div_n,exclude=channel)

        # move channel to the end
        x_scaled = np.moveaxis(x_scaled, channel, -1)
        channel = -1

        # u1: first rotation and prediction
        x_rot1   = self._rotate(x_scaled, axis=1, copy=False)
        u_rot1   = predict_direct(self.keras_model, x_rot1, channel_in=channel, channel_out=channel, single_sample=False,
                                  batch_size=batch_size, verbose=0)
        u1       = self._rotate(u_rot1, -1, axis=1, copy=False)

        # u2: second rotation and prediction
        x_rot2   = self._rotate(self._rotate(x_scaled, axis=2, copy=False), axis=0, copy=False)
        u_rot2   = predict_direct(self.keras_model, x_rot2, channel_in=channel, channel_out=channel, single_sample=False,
                                  batch_size=batch_size, verbose=0)
        u2       = self._rotate(self._rotate(u_rot2, -1, axis=0, copy=False), -1, axis=2, copy=False)

        n_channel_predicted = self.config.n_channel_out * (2 if self.config.probabilistic else 1)
        u_rot1.shape[channel] == n_channel_predicted or _raise(ValueError())
        u_rot2.shape[channel] == n_channel_predicted or _raise(ValueError())

        # move channel back to the front
        u1 = np.moveaxis(u1, channel, 0)
        u2 = np.moveaxis(u2, channel, 0)