pyTrainSeg/pyTMLS_4D_test.ipynb

Part 1: feature stack creation and saving

In [1]:

# necessary packages for feature stack
import numpy as np
from scipy import ndimage

import matplotlib.pyplot as plt

from skimage import filters, feature, io
from skimage.morphology import disk,ball

# from sklearn.ensemble import RandomForestClassifier

import os
import imageio
import sys
import dask
import dask.array
# import cupy as cp
# import cucim
from itertools import combinations_with_replacement
import xarray as xr

1.1 feature functions

In [2]:

# functions take chunked dask-array as input
def nd_gaussian(da, sig = 0):
    if np.abs(sig-0)<0.1:
        G = np.array(da)
        fullname = 'original'
    else:
        deptharray = np.ones(da.ndim)+4*sig
        deptharray = tuple(np.min([deptharray, da.shape], axis=0))
        G = da.map_overlap(filters.gaussian, depth=deptharray, boundary='nearest', sigma = sig).compute()
        # G = da.map_overlap(filters.gaussian, depth=4*sig+1, boundary='none', sigma = sig).compute()
        fullname = ''.join(['gaussian_',f'{sig:.1f}'])
    return G, fullname

#TODO create a class that makes the feature stacks
def nd_gaussian_stack(da, sigmas):
    fullnames = []
    gstack = np.zeros(list(da.shape) + [len(sigmas)])
    for sig,i in zip(sigmas, range(len(sigmas))):
        gstack[...,i], name = nd_gaussian(da, sig)
        fullnames.append(name)
    return gstack, fullnames

In [3]:

def nd_diff_of_gaussian(gstack, sigmas):
#     #creates a stack of {size} (see below)
    n = len(sigmas)
    size = int(n*(n-1)/2)
    dstack = np.zeros(list(da.shape) + [size])
    fullnames = []
    cc = 0
    for i in range(1,n):
        for j in range(i):
            dstack[...,cc] = gstack[...,i] - gstack[...,j]
            name = ''.join(['diff_of_gauss_',f'{sigmas[i]:.1f}','_',f'{sigmas[j]:.1f}'])
            fullnames.append(name)
            cc = cc + 1
    return dstack, fullnames

In [15]:

def ball_4d(sig):
    bnd = np.zeros((sig*2+1,sig*2+1,sig*2+1,sig*2+1), dtype = bool)
    bnd[sig,sig,sig,sig] = True
    ecd = ndimage.distance_transform_edt(~bnd)
    bnd = (ecd<sig+0.01).astype(int)
    return bnd

def nd_rank_like_filter(da, sigma, option):
    """
     input
     da - chunked das array up to 4D
     sigma - kernel size, scalar
     option, str ('minimum', 'maximum', 'median')
    """
    if da.ndim == 2:
        fp = disk(sigma)
    if da.ndim == 3:
        fp = ball(sigma)
    if da.ndim == 4:
        fp = ball_4d(sigma)
        
    if option == 'minimum':
        fun = ndimage.minimum_filter
    elif option == 'maximum':
        fun = ndimage.maximum_filter
    elif option == 'median':
        fun = ndimage.median_filter
    else:
        print(option+' not available')
    deptharray = np.ones(da.ndim)+sigma
    deptharray = tuple(np.min([deptharray, da.shape], axis=0))
    M = da.map_overlap(fun, depth=deptharray, footprint=fp).compute()
    # M = da.map_overlap(fun, depth=sigma+1, footprint=fp).compute()
    fullname = ''.join([option,'_',f'{sigma:.1f}'])
    return M, fullname

def nd_rank_like_stack(da, sigmas, option):
    fullnames = []
    mstack = np.zeros(list(da.shape) + [len(sigmas)-1])
    for sig,i in zip(sigmas[1:], range(len(sigmas)-1)):
        mstack[...,i], name = nd_rank_like_filter(da, sig, option)
        fullnames.append(name)
    return mstack, fullnames   
    

In [5]:

def nd_Hessian_matrix(G):
    """
    copied from skimage.feature.hessian_matrix
    just directly using Gaussian fitered arrays and dask
    """
    
    daG = dask.array.from_array(G)
    gradients = dask.array.gradient(daG)
    axes = range(G.ndim)
    gradients = [gradients[ax0].compute() for ax0 in axes]
    H_elems = [dask.array.gradient(gradients[ax0], axis=ax1).compute() for ax0, ax1 in combinations_with_replacement(axes, 2)]
    elems = [(ax0,ax1) for ax0, ax1 in combinations_with_replacement(axes, 2)]
    return H_elems, elems, gradients

def nd_Hessian_stack(G, sigma):
    H_elems, elems, gradients = nd_Hessian_matrix(G)
    hstack = np.zeros(list(G.shape)+[len(elems)])
    gradstack = np.zeros(list(G.shape)+[len(gradients)])
    
    #TODO: this is slow, find some better numpy function
    gradnames = []
    for i in range(len(elems)):
        hstack[...,i] = H_elems[i]
    for i in range(len(gradients)):
        gradnames.append(''.join(['gradient_',str(i),'_',f'{sigma:.1f}']))
        gradstack[...,i] = gradients[i]
    
    # print('got Hessian matrices, now doing the eigs')
    # eigs = feature.hessian_matrix_eigvals(H_elems) 
 # for now ignore the eigenvalues (too computationally expensive and H_elems already contains the image curvature  

    fullnames = []
    for i,j in elems:
        fullname = ''.join(['hessian_',str(i),str(j),'_',f'{sigma:.1f}'])
        fullnames.append(fullname) 
    
    stack = np.concatenate([hstack,gradstack], axis=-1)
    fullnames = fullnames+gradnames                         
        
    return stack, fullnames

def nd_Hessian_stacks(gstack, sigmas):
    flag = True
    fullnames = []
    for (i, sigma) in zip(range(gstack.shape[-1]), sigmas):
        a, b = nd_Hessian_stack(gstack[...,i], sigma)
        asize = a.shape[-1]
        if flag:
            flag = False
            hstacks = np.zeros(list(gstack[...,-1].shape)+[len(sigmas)*asize])
        hstacks[...,i*asize:i*asize+asize] = a
        fullnames = fullnames + b
    return hstacks, fullnames

In [6]:

def nd_feature_Stack(da, sigmas, feat_select):
#     TODO: make more elegant
    
    fstack = []
    featnames = []
    
    print('apply Gaussian filters anyway')
    gstack, gnames = nd_gaussian_stack(da, sigmas)
    
    if feat_select['Gaussian']:
        featnames = featnames + gnames
        fstack.append(gstack)
        
    if feat_select['Hessian']:
        print('get Hessian matrices')
        hstack, hnames = nd_Hessian_stacks(gstack, sigmas)
        featnames = featnames + hnames 
        fstack.append(hstack)
        
    if feat_select['Diff of Gaussians']:
        print('get differences of Gaussians')
        dstack, dnames = nd_diff_of_gaussian(gstack, sigmas)
        featnames = featnames + dnames
        fstack.append(dstack)
        
    if feat_select['maximum']:
        print('apply maximum filters')
        maxstack, maxnames = nd_rank_like_stack(da, sigmas, option='maximum')
        featnames = featnames + maxnames
        fstack.append(maxstack)
        
    if feat_select['median']:
        print('apply median filters')
        medstack, mednames = nd_rank_like_stack(da, sigmas, option='median')
        featnames = featnames + mednames
        fstack.append(medstack)
        
    if feat_select['minimum']:
        print('apply minimum filters')
        minstack, minnames = nd_rank_like_stack(da, sigmas, option='minimum')
        featnames = featnames + minnames
        fstack.append(minstack)
        
    return np.concatenate(fstack, axis=-1), featnames

def feat_stack_to_nc(fstack, featnames, path = None, name = 'test'):
    #TODO: include metadata and make general, now only 4D possible
    data = xr.Dataset({'feature_stack': (['x','y','z','time', 'feature'], fstack)},
                       coords = {'x': np.arange(fstack.shape[0]),
                                 'y': np.arange(fstack.shape[1]),
                                 'z': np.arange(fstack.shape[2]),
                                 'time': np.arange(fstack.shape[3]),
                                 'feature': featnames}
                      # attrs = feat_select
                     )
    # data.attrs['name'] = name
    if path is not None:
        data.to_netcdf(path)
    return data
    

1.2 load 4D dataset into dask-array TODO: employ datastreaming for datasets that are too large

In [3]:

#TODO
da = None

In [32]:

### provisoric manual loading of Jerermy's data
step1 = io.imread('/home/fische_r/NAS/testing/Jeremy_data/cropped/step_01.tif')

In [34]:

step1.shape

Out [34]:

(1500, 39, 233)

In [37]:

J_4D_data = np.zeros(list(step1.shape)+[8])

In [38]:

folder = '/home/fische_r/NAS/testing/Jeremy_data/cropped'
files = os.listdir(folder)
files.sort()
for i in range(len(files)):
    print(files[i])
    J_4D_data[...,i] = io.imread(os.path.join(folder, files[i]))

step_01.tif
step_02.tif
step_03.tif
step_04.tif
step_05.tif
step_06.tif
step_07.tif
step_08.tif

In [40]:

plt.imshow(J_4D_data[500,:,:,4])

Out [40]:

<matplotlib.image.AxesImage at 0x7f8e6bf738b0>

In [42]:

step1=None
J_4D_data = J_4D_data[500:700,...]

In [43]:

data = xr.Dataset({'tomo': (['x','y','z','t'], J_4D_data)},
                  coords = {'x': np.arange(J_4D_data.shape[0]),
                            'y': np.arange(J_4D_data.shape[1]),
                            'z': np.arange(J_4D_data.shape[2]),
                            't': np.arange(J_4D_data.shape[3])})

In [44]:

data.to_netcdf(os.path.join(folder, 'tomodata.nc'))

In [7]:

datapath = '/mpc/homes/fische_r/testing/tomodata.nc'
rawdata = xr.load_dataset(datapath)

In [8]:

da = dask.array.from_array(rawdata['tomo'].data)

In [9]:

da

Out [9]:

Array Chunk Bytes 110.93 MiB 110.93 MiB Shape (200, 39, 233, 8) (200, 39, 233, 8) Count 1 Tasks 1 Chunks Type float64 numpy.ndarray

1.3 feature settings

In [10]:

sigmas = [0, 1,4, 8]  #hard-coded for now, sobel and hessian require that first sigma is 0, diff, gaussian(sig=0) = 0

# default feature choice
feat_select = {'Gaussian': True, 
               # 'Sobel': True,
               'Hessian': True,
               'Diff of Gaussians': True,
               'maximum': True,
               'minimum': True,
               'median': True
              }

In [11]:

training_dict = {}

1.4 Create and store feature stack

In [12]:

#output path
outpath = '/mpc/homes/fische_r/testing/features.nc'

In [13]:

feat_data.close()

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Input In [13], in <cell line: 1>()
----> 1 feat_data.close()

NameError: name 'feat_data' is not defined

In [16]:

#this may take a while depending on the size and system
# may even crash because out of memory

fstack, featnames = nd_feature_Stack(da, sigmas, feat_select)

apply Gaussian filters anyway
get Hessian matrices
get differences of Gaussians
apply maximum filters
apply median filters
apply minimum filters

In [381]:

feat_data = feat_stack_to_nc(fstack, featnames, outpath, feat_select)

In [382]:

fstack = None
data = None

Part 2: Training the classifier

In [36]:

# necessary packages

#reload after kernel reset
import xarray as xr
import os
from skimage import io
import matplotlib.pyplot as plt
import numpy as np

#the classifier
from sklearn.ensemble import RandomForestClassifier
#stuff for painting on the image
from ipywidgets import Image
from ipywidgets import ColorPicker, IntSlider, link, AppLayout, HBox
from ipycanvas import  hold_canvas,  MultiCanvas #RoughCanvas,Canvas,

2.1 open (not load) feature stack

In [612]:

featpath =  '/mpc/homes/fische_r/testing/features.nc'
trainingpath = '/mpc/homes/fische_r/testing/'
# feat_data = xr.open_dataset(featpath)
feat_names = feat_data['feature'].data

2.2 training functions

In [38]:

def extract_training_data(im, truth, feat_stack):
    #pixelwise training data
    phase1 = truth==1
    phase2 = truth==2
    phase3 = truth==4   
    X1 = feat_stack[phase1]
    y1 = np.zeros(X1.shape[0])
    X2 = feat_stack[phase2]
    y2 = np.ones(X2.shape[0])
    X3 = feat_stack[phase3]
    y3 = 2*np.ones(X3.shape[0])

    y = np.concatenate([y1,y2,y3])
    X = np.concatenate([X1,X2,X3])
    
    return X,y

In [39]:

 def classify(X,y,im, feat_stack):
    # TODO: allow choice and manipulation of ML method
    clf =  RandomForestClassifier(n_estimators = 300, n_jobs=-1, random_state = 42, max_features=None) 
    clf.fit(X, y)
    num_feat = feat_stack.shape[2]
    ypred = clf.predict(feat_stack.reshape(-1,num_feat))
    result = ypred.reshape(im.shape).astype(np.uint8)
    return result, clf

In [84]:

def training_function(im, truth, feat_stack, training_dict, slice_name):
    flag = False
    slices = list(training_dict.keys())
    if slice_name in slices: 
        slices.remove(slice_name)
    if len(slices)>0:
        flag = True
        Xall = training_dict[slices[0]][0]
        yall = training_dict[slices[0]][1]
        for i in range(1,len(slices)):
            Xall = np.concatenate([Xall, training_dict[slices[i]][0]])
            yall = np.concatenate([yall, training_dict[slices[i]][1]])
            
    X,y = extract_training_data(im, truth, feat_stack)
    
    print('training and classifying')
    
    if flag:
        Xt = np.concatenate([Xall,X])
        yt = np.concatenate([yall,y])
        Xall = None
        yall = None
    else:
        Xt = X
        yt = y  
    result, clf = classify(Xt, yt, im, feat_stack)
    
    # store training data of current slice in dict
    training_dict[slice_name] = (X,y)
    return result, clf, training_dict
    

2.3 select 2D-slice and create a training set

In [43]:

# randomly suggest slice for training
# TODO: take coordinates from feat_data
dimensions = ['x','y','z','time']

test_dims = np.random.choice(dimensions, 2, replace=False)
p1 = np.random.choice(range(len(feat_data[test_dims[0]])))
p2 = np.random.choice(range(len(feat_data[test_dims[1]])))

print('You could try ',test_dims[0],'=',p1,' and ',test_dims[1],'=',p2)
# ts = np.random.choice(range(num_ts))+1
# print('try time step ',ts )

You could try  x = 113  and  time = 1

In [600]:

#select 2D slice orthogonal to these two axes
#replace dimensions in .sel accordingly, d1c and d2c are needed for filename
#TODO: make more elegant
d1c = 'z'
d2c = 'time'
p1c = 150
p2c = 7
im = feat_data['feature_stack'].sel(z = p1c, time = p2c, feature='original').data
im8 = np.uint8(im)
feat_stack = feat_data['feature_stack'].sel(z = p1c, time = p2c).data

In [601]:

im8 = im-im.min()
im8 = im8/im8.max()*255
im8old = im8.copy()
im8 = im8-50
im8[im8<0]=0
im8=im8*255/150
im8[im8>255] = 255

In [602]:

slice_name = ''.join([str(d1c),str(p1c),str(d2c),str(p2c)])
truthpath = os.path.join(trainingpath, ''.join(['label_image_',slice_name,'.tif']))

resultim = np.zeros(im.shape, dtype=np.uint8)
if os.path.exists(truthpath):
    truth = io.imread(truthpath)
    print('existing label set loaded')
else:
    truth = resultim.copy()

2.4 training interface

currently only 3 phases possible
manually specify #ff000 (phase 1) #00ff00 (phase 2) or #0000ff (phase 3) as color and paint

2.4.1 label data

In [609]:

width = im8.shape[1]
height = im8.shape[0]

Mcanvas = MultiCanvas(4, width=width, height=height)
background = Mcanvas[0]
resultdisplay = Mcanvas[2]
truthdisplay = Mcanvas[1]
canvas = Mcanvas[3]
canvas.sync_image_data = True

drawing = False
position = None
shape = []

def on_mouse_down(x, y):
    global drawing
    global position
    global shape

    drawing = True
    position = (x, y)
    shape = [position]

def on_mouse_move(x, y):
    global drawing
    global position
    global shape

    if not drawing:
        return

    with hold_canvas():
        canvas.stroke_line(position[0], position[1], x, y)

        position = (x, y)

    shape.append(position)

def on_mouse_up(x, y):
    global drawing
    global position
    global shape

    drawing = False

    with hold_canvas():
        canvas.stroke_line(position[0], position[1], x, y)
        canvas.fill_polygon(shape)

    shape = []

image_data = np.stack((im8, im8, im8), axis=2)
background.put_image_data(image_data, 0, 0)

# image_data = np.stack((imdry, imdry, imdry), axis=2)
# background.put_image_data(image_data, width, 0)

resultdisplay.global_alpha = 0.35
if np.any(resultim>0):
    result_data = np.stack((255*(resultim==0), 255*(resultim==1), 255*(resultim==2)), axis=2)
else:
    result_data = np.stack((0*resultim, 0*resultim, 0*resultim), axis=2)
resultdisplay.put_image_data(result_data, 0, 0)

canvas.on_mouse_down(on_mouse_down)
canvas.on_mouse_move(on_mouse_move)
canvas.on_mouse_up(on_mouse_up)

# canvas.stroke_style = "#749cb8"
# canvas.global_alpha = 0.75

picker = ColorPicker(description="Color:", value="#ff0000")
slidealpha = IntSlider(description="Result overlay", value=0.15)

link((picker, "value"), (canvas, "stroke_style"))
link((picker, "value"), (canvas, "fill_style"))
# link((slidealpha, "value"), (resultdisplay, "global_alpha"))

HBox((Mcanvas, picker, slidealpha))
#print('paint image with #ff0000 for air, #00ff00 for water and #0000ff for fiber')

HBox(children=(MultiCanvas(height=200, width=39), ColorPicker(value='#ff0000', description='Color:'), IntSlide…

In [610]:

diff = im8-imdry
diff[diff<0] = 0
fig, axes = plt.subplots(1,4)
axes[0].imshow(resultim==1)
axes[1].imshow(diff, 'gray')
axes[2].imshow(imdry, 'gray')
axes[3].imshow(im8, 'gray')

Out [610]:

<matplotlib.image.AxesImage at 0x7fbf34a81b80>

In [607]:

#create truth image from image, save to file
label_set = canvas.get_image_data()

truth[label_set[:,:,0]>0] = 1
truth[label_set[:,:,1]>0] = 2
truth[label_set[:,:,2]>0] = 4
imageio.imsave(truthpath, truth)

2.4.2 actual training

In [608]:

resultim, clf, training_dict = training_function(im, truth, feat_stack, training_dict, slice_name)

training and classifying

2.4.3 iterative training: go back to 2.4.1 until good segmentation or try different slice at 2.3

2.4.4 plot classifier properties

In [613]:

plt.figure( figsize=(16,9))
plt.stem(feat_names,clf.feature_importances_,'x')
plt.xticks(rotation=90)
plt.ylabel('importance')

Out [613]:

/tmp/ipykernel_268297/1483009324.py:2: MatplotlibDeprecationWarning: Passing the linefmt parameter positionally is deprecated since Matplotlib 3.5; the parameter will become keyword-only two minor releases later.
  plt.stem(feat_names,clf.feature_importances_,'x')

Text(0, 0.5, 'importance')

3 apply classifier to full data set

In [23]:

#TODO

In [470]:

training_dict = {}

In [25]:

da.ndim

Out [25]:

4

In [26]:

da.shape

Out [26]:

(200, 39, 233, 8)

In [28]:

deptharray = np.ones(da.ndim)+4
deptharray = np.max([deptharray, da.shape], axis=0)

In [31]:

tuple(deptharray)

Out [31]:

(200.0, 39.0, 233.0, 8.0)

In [ ]: