How to train the Unet for three-dimensional matrices （custom dataset） about tf_unet HOT 4 CLOSED

from tf_unet import unet
from tf_unet import util
from tf_unet import image_util
import numpy as np
import nibabel as nb
import matplotlib.pyplot as plt

# Read nii data
TrainFilePath = "G:/Tensorflow/tf_unet-master/demo/"
# TrainFilePath = "X:/tf_unet-master/demo/"
Train_Case = 1
Train_nii = nb.load(TrainFilePath + 'data-' + str(Train_Case) + '.nii')
Train_Label_nii = nb.load(TrainFilePath + 'data-label-' + str(Train_Case) + '.nii') # read label
# get data
Train_data = Train_nii.get_data()
Train_Label_data = Train_Label_nii.get_data() # 0: background; 1&2 ground truth
Train_data = Train_data.astype(np.uint8)
Train_Label_data = Train_Label_data.astype(np.uint8)
Train_Label_data[Train_Label_data == 2] = 1 # combine label 2 with label 1
# reshape data and label to [n, X, Y, channels], label [n, X, Y, classes].
nx = Train_data.shape[0]
ny = Train_data.shape[1]
nz = Train_data.shape[2]

Train_data = Train_data.reshape(nx, ny, nz, 1)
Train_data_label = Train_Label_data.reshape(nx, ny, nz, 1)
Train_data = Train_data.transpose(2, 0, 1, 3)
Train_data_label = Train_data_label.transpose(2, 0, 1, 3)
# provide data for the unet
data_provider = image_util.SimpleDataProvider(Train_data, Train_data_label)
# set parameter
net = unet.Unet(channels=1, n_class=3, layers=3, features_root=16)
trainer = unet.Trainer(net, optimizer="momentum", opt_kwargs=dict(momentum=0.2))
# train
path = trainer.train(data_provider, "./seg_3dim_data_trained_0509", training_iters=5, epochs=1, display_step=2) 
# predict
Test_Slice = Train_data[40]
Test_Slice = Test_Slice.reshape(1, nx, ny, 1)
Test_label = Train_data_label[40]
Test_label = Test_label.reshape(1, nx, ny, 1)

prediction = net.predict("./seg_3dim_data_trained_0509", Test_Slice)

# show the result
fig, ax = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(12,6))
ax[0].imshow(Test_Slice[0,...,0], cmap = 'gray', aspect="auto")
ax[1].imshow(Test_label[0,...,0], cmap = 'gray', aspect="auto")
mask = prediction[0, ..., 1] > 0.9
ax[2].imshow(mask, aspect="auto")

ax[0].set_title("Input")
ax[1].set_title("Ground truth")
ax[2].set_title("Prediction")
fig.tight_layout()
plt.show()

Traceback (most recent call last):
  File "E:\Program Files (x86)\JetBrains\PyCharm Community Edition 2017.1\helpers\pydev\pydev_run_in_console.py", line 78, in <module>
    globals = run_file(file, None, None)
  File "E:\Program Files (x86)\JetBrains\PyCharm Community Edition 2017.1\helpers\pydev\pydev_run_in_console.py", line 35, in run_file
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "E:\Program Files (x86)\JetBrains\PyCharm Community Edition 2017.1\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "G:/Tensorflow/tf_unet-master/Seg_3_dim_Nii_data.py", line 36, in <module>
    path = trainer.train(data_provider, "./seg_3dim_data_trained_0509", training_iters=5, epochs=1, display_step=2) 
  File "G:/Tensorflow/tf_unet-master\tf_unet\unet.py", line 403, in train
    test_x, test_y = data_provider(self.verification_batch_size)
  File "G:/Tensorflow/tf_unet-master\tf_unet\image_util.py", line 98, in __call__
    train_data, labels = self._load_data_and_label()
  File "G:/Tensorflow/tf_unet-master\tf_unet\image_util.py", line 58, in _load_data_and_label
    return train_data.reshape(1, ny, nx, self.channels), labels.reshape(1, ny, nx, self.n_class),
ValueError: cannot reshape array of size 262144 into shape (1,512,512,2)