class_docstr: This class is intended for the Synb0 model.
  The model is the deep learning part of the Synb0-Disco
  pipeline, thus stand-alone usage is not
  recommended.
docstr: Wrapper function to facilitate prediction of larger dataset.
  The function will pad the data to meet the required shape of image.
  Note that the b0 and T1 image should have the same shape
  
  Parameters
  ----------
  b0 : np.ndarray (batch, 77, 91, 77) or (77, 91, 77)
      For a single image, input should be a 3D array. If multiple images,
      there should also be a batch dimension.
  
  T1 : np.ndarray (batch, 77, 91, 77) or (77, 91, 77)
      For a single image, input should be a 3D array. If multiple images,
      there should also be a batch dimension.
  
  batch_size : int
      Number of images per prediction pass. Only available if data
      is provided with a batch dimension.
      Consider lowering it if you get an out of memory error.
      Increase it if you want it to be faster and have a lot of data.
      If None, batch_size will be set to 1 if the provided image
      has a batch dimension.
      Default is None
  
  average : bool
      Whether the function follows the Synb0-Disco pipeline and
      averages the prediction of 5 different models.
      If False, it uses the loaded weights for prediction.
      Default is True.
  Returns
  -------
  pred_output : np.ndarray (...) or (batch, ...)
      Reconstructed b-inf image(s)
docstr_tok: ["Wrapper","function","to","facilitate","prediction","of","larger","dataset",".","The","function","will","pad","the","data","to","meet","the","required","shape","of","image",".","Note","that","the","b0","and","T1","image","should","have","the","same","shape","Parameters","--","--","--","--","--","b0",":","np.ndarray","(","batch",",","77",",","91",",","77",")","or","(","77",",","91",",","77",")","For","a","single","image",",","input","should","be","a","3D","array",".","If","multiple","images",",","there","should","also","be","a","batch","dimension",".","T1",":","np.ndarray","(","batch",",","77",",","91",",","77",")","or","(","77",",","91",",","77",")","For","a","single","image",",","input","should","be","a","3D","array",".","If","multiple","images",",","there","should","also","be","a","batch","dimension",".","batch_size",":","int","Number","of","images","per","prediction","pass",".","Only","available","if","data","is","provided","with","a","batch","dimension",".","Consider","lowering","it","if","you","get","an","out","of","memory","error",".","Increase","it","if","you","want","it","to","be","faster","and","have","a","lot","of","data",".","If","None",",","batch_size","will","be","set","to","1","if","the","provided","image","has","a","batch","dimension",".","Default","is","None","average",":","bool","Whether","the","function","follows","the","Synb0-Disco","pipeline","and","averages","the","prediction","of","5","different","models",".","If","False",",","it","uses","the","loaded","weights","for","prediction",".","Default","is","True",".","Returns","--","--","--","-","pred_output",":","np.ndarray","(","...",")","or","(","batch",",","...",")","Reconstructed","b-inf","image","(","s",")"]
code: def predict(self, b0, T1, batch_size=None, average=True):
      r"""
      Wrapper function to facilitate prediction of larger dataset.
      The function will pad the data to meet the required shape of image.
      Note that the b0 and T1 image should have the same shape
  
      Parameters
      ----------
      b0 : np.ndarray (batch, 77, 91, 77) or (77, 91, 77)
          For a single image, input should be a 3D array. If multiple images,
          there should also be a batch dimension.
  
      T1 : np.ndarray (batch, 77, 91, 77) or (77, 91, 77)
          For a single image, input should be a 3D array. If multiple images,
          there should also be a batch dimension.
  
      batch_size : int
          Number of images per prediction pass. Only available if data
          is provided with a batch dimension.
          Consider lowering it if you get an out of memory error.
          Increase it if you want it to be faster and have a lot of data.
          If None, batch_size will be set to 1 if the provided image
          has a batch dimension.
          Default is None
  
      average : bool
          Whether the function follows the Synb0-Disco pipeline and
          averages the prediction of 5 different models.
          If False, it uses the loaded weights for prediction.
          Default is True.
      Returns
      -------
      pred_output : np.ndarray (...) or (batch, ...)
          Reconstructed b-inf image(s)
  
      """
      # Check if shape is as intended
      if (
          all([b0.shape[1:] != (77, 91, 77), b0.shape != (77, 91, 77)])
          or b0.shape != T1.shape
      ):
          raise ValueError(
              "Expected shape (batch, 77, 91, 77) or \
                               (77, 91, 77) for both inputs"
          )
  
      dim = len(b0.shape)
  
      # Add batch dimension if not provided
      if dim == 3:
          T1 = np.expand_dims(T1, 0)
          b0 = np.expand_dims(b0, 0)
      shape = b0.shape
  
      # Pad the data to match the model's input shape
      T1 = np.pad(T1, ((0, 0), (2, 1), (3, 2), (2, 1)), "constant")
      b0 = np.pad(b0, ((0, 0), (2, 1), (3, 2), (2, 1)), "constant")
  
      # Normalize the data.
      p99 = np.percentile(b0, 99, axis=(1, 2, 3))
      for i in range(shape[0]):
          T1[i] = normalize(T1[i], 0, 150, -1, 1)
          b0[i] = normalize(b0[i], 0, p99[i], -1, 1)
  
      if dim == 3:
          if batch_size is not None:
              logger.warning(
                  "Batch size specified, but not used",
                  "due to the input not having \
                                 a batch dimension",
              )
          batch_size = 1
  
      # Prediction stage
      if average:
          mean_pred = np.zeros(shape + (5,), dtype=np.float32)
          for i in range(5):
              self.fetch_default_weights(i)
              temp = np.stack([b0, T1], -1)
              input_data = np.moveaxis(temp, 3, 1).astype(np.float32)
              prediction = np.zeros(
                  (shape[0], 80, 80, 96, 1), dtype=np.float32
              )
              for batch_idx in range(
                  batch_size, shape[0] + 1, batch_size
              ):
                  temp_input = input_data[
                      batch_idx - batch_size : batch_idx
                  ]
                  temp_pred = self.__predict(temp_input)
                  prediction[
                      batch_idx - batch_size : batch_idx
                  ] = temp_pred
              remainder = np.mod(shape[0], batch_size)
              if remainder != 0:
                  temp_pred = self.__predict(input_data[-remainder:])
                  prediction[-remainder:] = temp_pred
              for j in range(shape[0]):
                  temp_pred = unnormalize(
                      prediction[j], -1, 1, 0, p99[j]
                  )
                  prediction[j] = temp_pred
  
              prediction = prediction[:, 2:-1, 2:-1, 3:-2, 0]
              prediction = np.moveaxis(prediction, 1, -1)
  
              mean_pred[..., i] = prediction
          prediction = np.mean(mean_pred, axis=-1)
      else:
          temp = np.stack([b0, T1], -1)
          input_data = np.moveaxis(temp, 3, 1).astype(np.float32)
          prediction = np.zeros(
              (shape[0], 80, 80, 96, 1), dtype=np.float32
          )
          for batch_idx in range(batch_size, shape[0] + 1, batch_size):
              temp_input = input_data[
                  batch_idx - batch_size : batch_idx
              ]
              temp_pred = self.__predict(temp_input)
              prediction[:batch_idx] = temp_pred
          remainder = np.mod(shape[0], batch_size)
          if remainder != 0:
              temp_pred = self.__predict(input_data[-remainder:])
              prediction[-remainder:] = temp_pred
          for j in range(shape[0]):
              prediction[j] = unnormalize(
                  prediction[j], -1, 1, 0, p99[j]
              )
  
          prediction = prediction[:, 2:-1, 2:-1, 3:-2, 0]
          prediction = np.moveaxis(prediction, 1, -1)
  
      if dim == 3:
          prediction = prediction[0]
  
      return prediction
  
imports: import logging
  import numpy
  from dipy.data import get_fnames
  from dipy.testing.decorators import doctest_skip_parser
  from dipy.utils.optpkg import optional_package
  from dipy.nn.utils import normalize, unnormalize, set_logger_level
label_desc: Use image node_id 5 for calling the Synb0 obj's underlying member method code with example usage: obj.predict(b0, T1, batch_size, average) and returns: prediction