demo_train.py

# -*- coding: utf-8 -*-
"""
Demo showing how to train stneuronet from selected data from Allen Brain Observatory dataset

@author: Somayyeh Soltanian-Zadeh
%
% Please cite this paper if you use any component of this software:
% S. Soltanian-Zadeh, K. Sahingur, S. Blau, Y. Gong, and S. Farsiu, "Fast 
% and robust active neuron segmentation in two-photon calcium imaging using 
% spatio-temporal deep learning," Proceedings of the National Academy of Sciences (PNAS), 2019.
%
% Released under a GPL v2 license.
"""
import os
import sys
import numpy as np
import math
from pathlib import Path
import scipy.io as sio

import niftynet

import STNeuroNetPkg
import matlab 
matlabLib = STNeuroNetPkg.initialize()

## Set directories
dirpath = os.getcwd()
DataDir = os.path.join(dirpath,'Dataset','ABO')
DirSaveData = os.path.join(dirpath,'Results','ABO','data')
DirSaveTempMask = os.path.join(dirpath,'Results','ABO','data','TempMask')
DirSave = os.path.join(dirpath,'Results','ABO','Probability map')
DirModel = os.path.join(dirpath,'models','ABO','Network','1')
DirMask = os.path.join(dirpath,'Markings','ABO','Layer275','Grader1')

## Set parameters
fs = matlab.double([6])              #recording speed after temporal downsampling, Hz
SZ = matlab.double([487,487])
fNeuron = matlab.double([2.9])           # firing rate of neurons, Hz
Pd = matlab.double([1-0.035])         #Probability of detection
tau = 0.2               # decay rate of calcium sensor, seconds
pixSize = 0.78         #um
meanR = 5.85           # neuron radius in um
AvgArea = round(math.pi*(meanR/pixSize)**2)
maxDist = 5/pixSize    # 5 um -> pixels
ds = 5
border = 13
s = 30
Wt = matlab.double([200])     #length of temporal window for subsampling training data
Ws = matlab.double([144])     # length of spatial window
flag = 0                      # If tag is one, then overlapping patches from the video will be saved
                              # in a directory to be used during training (along with their rotations).
                              # Make sure to have enough memory to save the data. If flag is zero, small
                              # patches will be randomly cropped during training on the fly. However
                              # if the data is very large, this will be very slow.

# Hyperparameters
thresh = matlab.double(np.concatenate((np.arange(.50,.95,.05),np.arange(.96,1,.02))).tolist())
minArea =  matlab.double(np.concatenate(([15,20],np.arange(25,100,15))).tolist())    #in pixels

## Preprocess Train data specified by name. Note that any data name
## in the data directory that is not included in "name" will be ignored in training
name = ['501484643']

NormVals = matlab.double([0,0])
#%
for ind in range(len(name)):
    data_file = Path(DataDir+'//'+ name[ind]+'_processed.nii.gz')
    if not data_file.exists():
        raise ValueError('Downsampled, cropped data file in .nii.gz format not found.')
    my_file = Path(DirSaveData+'//'+ name[ind]+'_dsCropped_HomoNorm.nii.gz')
    if not my_file.exists():
        data_file = os.path.join(DataDir, name[ind]+'_processed.nii.gz')
        matlabLib.HomoFilt_Normalize(data_file,DirSaveData,name[ind],s,nargout=0)
    my_file = Path(DirSaveTempMask+'//TemporalMask_'+ name[ind]+'.nii.gz')
    if not my_file.exists():
        if not os.path.isdir(DirSaveTempMask):
            os.mkdir(DirSaveTempMask)
        matlabLib.prepareTemporalMask(DataDir,DirMask,DirSaveTempMask,name[ind],tau,fs,fNeuron,Pd,nargout=0)
    if flag:
        TrainDirData = os.path.join(dirpath,'Results','ABO','subImages')
        TrainDirMask = TrainDirData
        if not os.path.isdir(TrainDirData):
            os.mkdir(TrainDirData)
        matlabLib.SaveSubVolumes(DirSaveData,DirSaveTempMask,TrainDirData,name[ind],Ws,Wt,nargout=0)
        Nmask = 'TempMask_'     #train Mask names contains
        Nvid = 'HomoVid_'       #train video names contains
    else:
         TrainDirData = DirSaveData
         TrainDirMask = DirSaveTempMask
         Nmask = 'TemporalMask_'
         Nvid = '_dsCropped_HomoNorm'

#%%
## train the network
# first create a new config file based on the current data
f = open("demo_config_empty.ini")
mylist = f.readlines()
f.close()

mylistInf = mylist

indPath = []
indName = []
indNoName = []
indModel = []
indSave = []
for ind in range(len(mylist)):
    if mylist[ind].find('path_to_search')>-1:
        indPath.append(ind)
    if mylist[ind].find('filename_contains')>-1:
        indName.append(ind)
    if mylist[ind].find('filename_not_contains')>-1:
        indNoName.append(ind)          
    if mylist[ind].find('model_dir')>-1:
        indModel.append(ind) 
    if mylist[ind].find('save_seg_dir')>-1:
        indSave.append(ind)            
        
# write path of data
mystr = list(mylist[indPath[0]])
mystr = "".join(mystr[:-1]+ list(TrainDirData) + list('\n'))
mylist[indPath[0]] = mystr
# write name of data
mystr = list(mylist[indName[0]])
mystr = "".join(mystr[:-1]+ list(Nvid) + list('\n'))
mylist[indName[0]] = mystr

# write information for labels
mystr = list(mylist[indPath[1]])
mystr = "".join(mystr[:-1]+ list(TrainDirMask) + list('\n'))
mylist[indPath[1]] = mystr
# write name of data
mystr = list(mylist[indName[1]])
mystr = "".join(mystr[:-1]+ list(Nmask) + list('\n'))
mylist[indName[1]] = mystr

# exclude any other data not listed in names
AllFiles = os.listdir(TrainDirMask)
AllNames = []
for ind in range(len(AllFiles)):
    if AllFiles[ind].find(Nmask)>-1:
        ind1 =AllFiles[ind].find(Nmask)+len(Nmask)
        ind2 = ind1+9
        AllNames.append(AllFiles[ind][ind1:ind2])
        
excludeNames = [c for c in AllNames if c not in name]    
if len(excludeNames):   
    mystr = list(mylist[indNoName[0]])
    temp = mystr[:-1] + list(',')
    for ind in range(len(excludeNames)):
        temp = temp + list(excludeNames[ind]) + list(',')
    mystr = "".join(temp[:-1]+ list('\n'))
    mylist[indNoName[0]] = mystr
    mylist[indNoName[1]] = mystr

#write where model is located
mystr = list(mylist[indModel[0]])
mystr = "".join(mystr[:-1]+ list(DirModel) + list('\n'))
mylist[indModel[0]] = mystr


# Write to a new config file
f = open('config_train.ini','w')
f.write(''.join(mylist))
f.close()


sys.argv=['','train','-a','net_segment','--conf',os.path.join('config_train.ini'),'--batch_size','3']
niftynet.main()

#%%
# inference config file
f = open("demo_config_empty.ini")
mylistInf = f.readlines()
f.close()
# write path of data
mystr = list(mylistInf[indPath[0]])
mystr = "".join(mystr[:-1]+ list(DirSaveData) + list('\n'))
mylistInf[indPath[0]] = mystr
# write name of data
mystr = list(mylistInf[indName[0]])
mystr = "".join(mystr[:-1]+ list('_dsCropped_HomoNorm') + list('\n'))
mylistInf[indName[0]] = mystr

mylistInf[indNoName[0]] = mylist[indNoName[0]]
mylistInf[indModel[0]] = mylist[indModel[0]]
mylistInf[indNoName[1]] = mylist[indNoName[1]]

#write where model is located
mystr = list(mylist[indModel[0]])
mystr = "".join(mystr[:-1]+ list(DirModel) + list('\n'))
mylist[indModel[0]] = mystr

#write where to save result
mystr = list(mylistInf[indSave[0]])
mystr = "".join(mystr[:-1]+ list(DirSave) + list('\n'))
mylistInf[indSave[0]] = mystr

# Write to a new config file
f = open('config_inf.ini','w')
f.write(''.join(mylistInf))
f.close()

## Apply the trained network to train data for further processing
sys.argv=['','inference','-a','net_segment','--conf','config_inf.ini','--batch_size','1']
niftynet.main()

#%%
## Postprocess to get individual neurons
F1All = np.zeros([len(name),thresh.size[-1],minArea.size[-1]])
if DirMask is not None:
    for ind in range(len(name)):
        recall,precision,F1 = matlabLib.multiple_postProcess(SZ,DirSave,DirMask,name[ind],AvgArea,minArea,thresh,0.5,nargout=3)
        F1All[ind,:,:] = np.array(F1)

    # Compute average F1 across data to get best hyperparameters            
    F1mean = np.mean(F1All,axis=0)  
    ind = np.argmax(F1mean.flatten())
    indA = math.floor(ind/minArea.size[-1])
    indT = ind - minArea.size[-1]*indA 
    
    print('Best thresh: %f , minArea: %f:'%(np.array(thresh)[:,indT],np.array(minArea)[:,indA]))

#%%
matlabLib.terminate()