# Path: filtering/riveal.py
# def riveal(img, swc, K=9, nsample=8e4, epoch=20):
#     print('-- oiginal image size: ', img.shape)
#     K = math.floor(K)  # Make sure K is integer to avoid warnings
#     # Pad the image and swc
#     margin = 3 * K
#     img = padimg(img, margin)
#     swc = padswc(swc, margin)
# 
#     # Make the skeleton distance transform
#     print('-- Distance transform for swc...')
#     dt, foreground_region = make_skdt(img.shape, swc, K)
# 
#     # Normalise data
#     # img = standardise(img)
#     # dt = standardise(dt)
#     img /= img.max()
#     # dt /= dt.max()
# 
#     # Make the confident region
#     print('==swc shape:', swc.shape)
#     print('-- Making the confidence regions...(1/4)')
#     high_conf_region = make_conf_region(
#         img.shape, swc, K, low_conf=0.5, high_conf=1.)
#     # print('-- Making the confidence regions...(2/4)')
#     # mid_conf_region = make_conf_region(img.shape, swc, K,
#     #                                    low_conf=0.25, high_conf=0.5)
#     print('-- Making the confidence regions...(3/4)')
#     low_conf_region = make_conf_region(
#         img.shape, swc, K, low_conf=0., high_conf=0.25)
# 
#     # # Fill only the central part of background region
#     print('-- Making the confidence regions...(4/4)')
#     background_region = np.zeros(img.shape)
#     bg = np.logical_not(foreground_region)
#     bg = np.logical_and(foreground_region, img > 0)
#     for i in range(3):
#         bg = binary_dilation(bg)
#     background_region[margin:-margin, margin:-margin, margin:-margin] = bg[
#         margin:-margin, margin:-margin, margin:-margin]
# 
#     from matplotlib import pyplot as plt
#     plt.subplot(3, 1, 1)
#     plt.imshow(high_conf_region.max(axis=-1))
#     plt.title('high conf')
#     plt.subplot(3, 1, 2)
#     plt.imshow(low_conf_region.max(axis=-1))
#     plt.title('low conf')
#     plt.subplot(3, 1, 3)
#     plt.imshow(background_region.max(axis=-1))
#     plt.title('bg')
#     plt.show()
# 
#     # Randomly sample 2.5D blocks from the include region
#     print('-- Sampling blocks')
#     x1, y1 = sample_block(img, dt, high_conf_region, K,
#                           math.ceil(nsample * 0.75))
#     # x2, y2 = sample_block(img, dt, mid_conf_region,
#     #                       K, math.ceil(nsample * 0.2))
#     x3, y3 = sample_block(img, dt, low_conf_region, K,
#                           math.ceil(nsample * 0.1))
#     y3.fill(0.)
#     x4, y4 = sample_block(img, dt, background_region, K,
#                           math.ceil(nsample * 0.15))
#     y4.fill(0.)
#     train_x = np.vstack((x1, x3, x4))
#     train_y = np.vstack((y1, y3, y4))
# 
#     # Build the CNN with keras+tensorflow
#     print('--Training CNN...')
#     model = traincnn(train_x, train_y, K, epoch)
# 
#     # Make the prediction within an area larger than
#     # the segmentation of the image
#     print('-- Predicting...')
#     bimg = img > 0
#     for i in range(6):
#         bimg = binary_dilation(bimg)
#     include_region = bimg > 0
#     include_idx = np.argwhere(include_region)
#     nidx = include_idx.shape[0]
# 
#     predict_x = np.zeros((nsample, 2 * K + 1, 2 * K + 1, 3))
#     rest = nidx
#     resultimg = np.zeros(img.shape)
#     pbar = tqdm(total=nidx)
# 
#     # Predict every batch of blocks
#     while rest > 0:
#         startidx = -rest
#         endidx = -rest + nsample if -rest + nsample < nidx else nidx
#         rest -= nsample
# 
#         # Write the value to each include voxel
#         for i, gidx in enumerate(range(int(startidx), int(endidx))):
#             bx, by, bz = include_idx[gidx, :]
#             predict_x[i, :, :, 0] = img[bx - K:bx + K + 1, by - K:by + K + 1,
#                                         bz]
#             predict_x[i, :, :, 1] = img[bx - K:bx + K + 1, by, bz - K:bz + K +
#                                         1]
#             predict_x[i, :, :, 2] = img[bx, by - K:by + K + 1, bz - K:bz + K +
#                                         1]
# 
#         pd = model.predict(predict_x, batch_size=64, verbose=0).flatten()
# 
#         for i, gidx in enumerate(range(int(startidx), int(endidx))):
#             bx, by, bz = include_idx[gidx, :]
#             resultimg[bx, by, bz] = pd[i]
#         pbar.update(nsample)
# 
#     resultimg = unpadimg(resultimg, margin)
# 
#     return resultimg
# 
# Path: filtering/thresholding.py
# def rescale(img, overwrite=False):
#     '''
#     Rescale image intensities linearly to 0~255
#     '''
#     if not overwrite:
#         result = img.copy()
#     else:
#         result = img
# 
#     result = result.astype('float')
#     result -= result.min()
#     result /= result.max()
#     result *= 255
#     return result


from filtering.riveal import riveal
from rivuletpy.utils.io import *
from filtering.thresholding import rescale
    from skimage import filters
    from skimage import filter as filters

img = loadimg('tests/data/test.tif')
dtype = img.dtype
swc = loadswc('tests/data/test.swc')
img = riveal(img, swc, nsample=5e4, epoch=30)

# Path: filtering/riveal.py
# def riveal(img, swc, K=9, nsample=8e4, epoch=20):
#     print('-- oiginal image size: ', img.shape)
#     K = math.floor(K)  # Make sure K is integer to avoid warnings
#     # Pad the image and swc
#     margin = 3 * K
#     img = padimg(img, margin)
#     swc = padswc(swc, margin)
# 
#     # Make the skeleton distance transform
#     print('-- Distance transform for swc...')
#     dt, foreground_region = make_skdt(img.shape, swc, K)
# 
#     # Normalise data
#     # img = standardise(img)
#     # dt = standardise(dt)
#     img /= img.max()
#     # dt /= dt.max()
# 
#     # Make the confident region
#     print('==swc shape:', swc.shape)
#     print('-- Making the confidence regions...(1/4)')
#     high_conf_region = make_conf_region(
#         img.shape, swc, K, low_conf=0.5, high_conf=1.)
#     # print('-- Making the confidence regions...(2/4)')
#     # mid_conf_region = make_conf_region(img.shape, swc, K,
#     #                                    low_conf=0.25, high_conf=0.5)
#     print('-- Making the confidence regions...(3/4)')
#     low_conf_region = make_conf_region(
#         img.shape, swc, K, low_conf=0., high_conf=0.25)
# 
#     # # Fill only the central part of background region
#     print('-- Making the confidence regions...(4/4)')
#     background_region = np.zeros(img.shape)
#     bg = np.logical_not(foreground_region)
#     bg = np.logical_and(foreground_region, img > 0)
#     for i in range(3):
#         bg = binary_dilation(bg)
#     background_region[margin:-margin, margin:-margin, margin:-margin] = bg[
#         margin:-margin, margin:-margin, margin:-margin]
# 
#     from matplotlib import pyplot as plt
#     plt.subplot(3, 1, 1)
#     plt.imshow(high_conf_region.max(axis=-1))
#     plt.title('high conf')
#     plt.subplot(3, 1, 2)
#     plt.imshow(low_conf_region.max(axis=-1))
#     plt.title('low conf')
#     plt.subplot(3, 1, 3)
#     plt.imshow(background_region.max(axis=-1))
#     plt.title('bg')
#     plt.show()
# 
#     # Randomly sample 2.5D blocks from the include region
#     print('-- Sampling blocks')
#     x1, y1 = sample_block(img, dt, high_conf_region, K,
#                           math.ceil(nsample * 0.75))
#     # x2, y2 = sample_block(img, dt, mid_conf_region,
#     #                       K, math.ceil(nsample * 0.2))
#     x3, y3 = sample_block(img, dt, low_conf_region, K,
#                           math.ceil(nsample * 0.1))
#     y3.fill(0.)
#     x4, y4 = sample_block(img, dt, background_region, K,
#                           math.ceil(nsample * 0.15))
#     y4.fill(0.)
#     train_x = np.vstack((x1, x3, x4))
#     train_y = np.vstack((y1, y3, y4))
# 
#     # Build the CNN with keras+tensorflow
#     print('--Training CNN...')
#     model = traincnn(train_x, train_y, K, epoch)
# 
#     # Make the prediction within an area larger than
#     # the segmentation of the image
#     print('-- Predicting...')
#     bimg = img > 0
#     for i in range(6):
#         bimg = binary_dilation(bimg)
#     include_region = bimg > 0
#     include_idx = np.argwhere(include_region)
#     nidx = include_idx.shape[0]
# 
#     predict_x = np.zeros((nsample, 2 * K + 1, 2 * K + 1, 3))
#     rest = nidx
#     resultimg = np.zeros(img.shape)
#     pbar = tqdm(total=nidx)
# 
#     # Predict every batch of blocks
#     while rest > 0:
#         startidx = -rest
#         endidx = -rest + nsample if -rest + nsample < nidx else nidx
#         rest -= nsample
# 
#         # Write the value to each include voxel
#         for i, gidx in enumerate(range(int(startidx), int(endidx))):
#             bx, by, bz = include_idx[gidx, :]
#             predict_x[i, :, :, 0] = img[bx - K:bx + K + 1, by - K:by + K + 1,
#                                         bz]
#             predict_x[i, :, :, 1] = img[bx - K:bx + K + 1, by, bz - K:bz + K +
#                                         1]
#             predict_x[i, :, :, 2] = img[bx, by - K:by + K + 1, bz - K:bz + K +
#                                         1]
# 
#         pd = model.predict(predict_x, batch_size=64, verbose=0).flatten()
# 
#         for i, gidx in enumerate(range(int(startidx), int(endidx))):
#             bx, by, bz = include_idx[gidx, :]
#             resultimg[bx, by, bz] = pd[i]
#         pbar.update(nsample)
# 
#     resultimg = unpadimg(resultimg, margin)
# 
#     return resultimg
# 
# Path: filtering/thresholding.py
# def rescale(img, overwrite=False):
#     '''
#     Rescale image intensities linearly to 0~255
#     '''
#     if not overwrite:
#         result = img.copy()
#     else:
#         result = img
# 
#     result = result.astype('float')
#     result -= result.min()
#     result /= result.max()
#     result *= 255
#     return result


# Path: tests/test_riveal.py
from filtering.riveal import riveal
from rivuletpy.utils.io import *
from filtering.thresholding import rescale
    from skimage import filters
    from skimage import filter as filters


img = loadimg('tests/data/test.tif')
dtype = img.dtype
swc = loadswc('tests/data/test.swc')
img = riveal(img, swc, nsample=5e4, epoch=30)

img = rescale(img)