CCSegProcess

#!/usr/bin/env python3

# command line processing

import os
import sys
import numpy
import getopt
import errno
import shutil

import CCSegPipeThickness
import CCSegPipeMidSag
import CCSegPipeSeg
import CCSegPipeSegToNative
import CCSegUtils

import CCSegSubjectSelection
import traceback

def criticalError(errorString):
    print("Error: " + str(errorString))
    quit()

def usage():
    print("Usage: " + sys.argv[0] + " [options] <input directory> <output directory>")
    print()
    print("DESCRIPTION")
    print()
    print("Front end script for CC Segmentation and Thickness profile pipeline, processing components")
    print()
    print("Performs the following steps:")
    print("\t1. Midsagittal plane extraction")
    print("\t2. CC segmentation")
    print("\t3. Thickness profile generation")
    print("\t4 (optional). Projection of segmentation and parcellations back to native space")
    print()
    print("The above steps may be executed with the following switches")
    print()
    print("PROCESSING DIRECTIVES (do options)")
    print("\t--do-midsag: Step 1")
    print("\t--do-seg: Step 2")
    print("\t--do-thickness: Step 3")
    print("\t--do-segtonative: Step 4")
    print("\t--do-all: Steps 1, 2, 3, and 4")
    print("\t--do-thicknessmanedit: Step 3, but only process subjects that have manually edited segmentations")
    print()
    print("OPTIONS")
    print("Segmentation (--do-seg) options:")
    print("\t--seg-no-lk: do not perform Lucas-Kanade tracking")
    print("\t--seg-choose-first: choose the first initialisation point for the segmentation output")
    print("\t--seg-no-auto: do not perform automatic segmentation, leave empty segmentation on midsagittal slice, use if all automatic options fail")
    print()
    print("Thickness profile (--do-thickness) options:")
    print("\t--num-thickness-nodes=(integer): the number of streamlines and thickness profiles to generate, 100 if not given")
    print()
    print("Segmentation to native (--do-segtonative) options:")
    print("\t--dilate-parasag: when performing the backprojection to native, dilate the segmentation and parcellations to the neighbouring parasagittal slices")
    print("\t--dilate-parasag-slices=(integer): when performing the backprojection to native, dilate the segmentation a given number of slices, implies --dilate-parasag")
    print()
    print("Misc options:")
    print("\t--redo-done: each processing step will not overwrite existing output, this forces existing output to be overwritten")
    print("\t--save-graphics: output graphics showing results of processing steps, look in subdirectories of output directory")
    CCSegSubjectSelection.printHelpText()

def main():
    
    longOptStrings = ['do-midsag', 'do-seg', 'do-thickness', 'do-segtonative', 'do-all', 'do-thicknessmanedit'] # do directives
    longOptStrings.extend(['skullstripped', 'midsag-no-brainmask']) # midsag options
    longOptStrings.extend(['seg-no-lk', 'seg-no-auto', 'seg-choose-first', 'seg-grid-initpoints']) # seg options
    longOptStrings.extend(['num-thickness-nodes=', 'thickness-no-reg']) # thickess options
    longOptStrings.extend(['dilate-parasag', 'dilate-parasag-slices=']) # segtonative options
    longOptStrings.extend(['long-subject-timepoint-file=']) # longitudinal options
    longOptStrings.extend(['redo-done', 'save-graphics']) # misc options
    longOptStrings.extend(CCSegSubjectSelection.getOptions())    
    opts, args = getopt.getopt(sys.argv[1:], "h", longOptStrings)
    
    # initialise options

    # mandatory arguments
    inputDirectory = None
    outputDirectory = None
    
    # "Mode options"
    doMidsag = False
    doSeg = False
    doSegToNative = False
    doThickness = False
    doThicknessManedit = False

    MSPMethod = 'symmetric'
    numThicknessNodes = 100
    redoDone = False
    doGraphics = False
    dilateParaSagSlices = 0
    thicknessNoReg = False
    useMNIBrainMask = True

    quitOnError = False

    segLK = True
    segNoAuto = False
    segChooseFirst = False
    segGridInitPoints = True

    skullStripped = False

    longSubjectTimepointFile = None
    # do the mandatory arguments first

    if len(args) != 2:
        print("The number of arguments must be 2")
        usage()
        exit()
    numpy.set_printoptions(precision = 3, formatter = {'all':lambda x: "%.3f" % x})

    inputDirectory = args[0]
    outputDirectory = args[1]
    
    #print "inputDirectory = " + inputDirectory
    #print "outputDirectory = " + outputDirectory

    for o, a in opts:
        if o == '-h':
            usage()
            quit()
        if o == '--num-thickness-nodes':
            try:
                numThicknessNodes = int(a)
            except ValueError:
                criticalError("the --num-thickness-nodes value was not an integer")
        if o == '--redo-done':
            redoDone = True
        if o == '--msp-method':
            print('--msp-method is deprecated, ignored')
        if o == '--do-midsag':
            doMidsag = True
        if o == '--do-seg':
            doSeg = True
        if o == '--do-segtonative':
            doSegToNative = True
        if o == '--do-thickness':
            doThickness = True
        if o == '--do-all':
            doMidsag = True
            doSeg = True
            doSegToNative = True
            doThickness = True
        if o == '--do-thicknessmanedit':
            doThicknessManedit = True
        if o == '--save-graphics':
            doGraphics = True
        if o == '--skullstripped':
            skullStripped = True
        if o == '--midsag-no-brainmask':
            useMNIBrainMask = False
        if o == '--seg-no-lk':
            segLK = False
        if o == '--seg-no-auto':
            segNoAuto = True
        if o == '--seg-choose-first':
            segChooseFirst = True
        if o == '--seg-grid-initpoints':
            segGridInitPoints = True
        if o == '--dilate-parasag':
            dilateParaSagSlices = 1
        if o == '--dilate-parasag-slices':
            try:
                dilateParaSagSlices = int(a)
            except ValueError:
                criticalError("the --dilate-parasag-slices value was not an integer")
        if o == '--thickness-no-reg':
            thicknessNoReg = True
        if o == '--long-subject-timepoint-file':
            longSubjectTimepointFile = a
    #for o, a in opts:

    # check to see if the user has given at least one of the MSP methods
    #if MSPMethod == "acpcdetect":
    #    print "Warning: the acpcdetect method does not yet give me backprojection to native space"
    
    MSPMethods = ['flirt', 'flirt_affine', 'acpcdetect', 'long', 'symmetric', 'none']
    MSPMethods = ['symmetric']

    if not MSPMethod in MSPMethods:
        criticalError("Only the MSP Method \"symmetric\" is supported now, the others have been deprecated")

    if numThicknessNodes <= 0:
        criticalError("The number of thickness nodes specified was not positive")
    
    if not os.path.isdir(inputDirectory):
        criticalError("The input directory specified was not found")
    
    if doMidsag == False and doSeg == False and doSegToNative == False and doThickness == False and doThicknessManedit == False:
        print("Error: No processing steps specified")
        usage()
        quit()
    
    if doMidsag == True and doSeg == False and doThickness == True:
        print("Warning: --do-midsag and --do-thickness were specified but the middle step --do-seg was not, setting --do-seg")
        doSeg = True
    
    if dilateParaSagSlices < 0:
        print("--dilate-parasag-slices must be a positive integer")
        usage()
        quit()

    actualInputFiles = CCSegSubjectSelection.selectSubjectsFromOpts(inputDirectory, opts = opts)
    if len(actualInputFiles) == 0:
        criticalError("No subjects to process")
    
    CCSegUtils.mkdirSafe(outputDirectory)

    #quit()
    #03_S_1122_sc_seg.hdf5
    #for curFile in inputNIFTIFiles:
    for curSubject in sorted(actualInputFiles):
    #for curFile in ['003_S_1122_sc']:    
    #for curFile in ['051_S_1338_sc']:    
        #print "Current subject: " + curFile
        outputBase = os.path.join(outputDirectory, curSubject)

        if doMidsag:
            if redoDone or not os.path.isfile(outputBase + "_midsag.hdf5"):
                #try:
                    print(curSubject + ": midsag")
                    CCSegPipeMidSag.midsagExtract(os.path.join(inputDirectory, curSubject), outputBase, MSPMethod, doGraphics = doGraphics, skullStripped = skullStripped, useMNIBrainMask = useMNIBrainMask)
                #except AssertionError as e:
                #    print "MidSag Error: " + str(e) + ", subject " + outputBase
                #    if quitOnError:
                #        quit()
                #    else:
                #        continue
                #except Exception as e:
                #    print "Midsag Error: " + str(e) + ", subject " + outputBase
                #    print traceback.print_exc(file=sys.stdout)
                #    continue
        if doSeg:
            #print outputBase + "_midsag.hdf5"
            if redoDone or not os.path.isfile(outputBase + "_seg.hdf5"):
                print(curSubject + ": seg")
                if os.path.isfile(outputBase + "_seg_manedit.hdf5"):
                    print("Found manedit: " + outputBase + "_seg_manedit.hdf5" + ", but redoing segmentation, deleting")
                    try:
                        os.remove(outputBase + "_seg_manedit.hdf5")
                    except Exception as e:
                        print("Couldnt delete manedit file: " + outputBase + "_seg_manedit.hdf5" + ", subject " + outputBase)
                        if quitOnError:
                            quit()
                        else:
                            continue
                try:
                    CCSegPipeSeg.segCC(outputBase, doGraphics = doGraphics, doLK = segLK, segNoAuto = segNoAuto, segChooseFirst = segChooseFirst, segGridInitPoints = segGridInitPoints)
    
                except AssertionError as e:
                    print("CCSeg Error: " + str(e) + ", subject " + outputBase)
                    
                    L = os.listdir(os.path.join(outputDirectory, 'seg'))
                    CCSegUtils.mkdirSafe(os.path.join(outputDirectory, 'seg_bad'))
                    for curFile in sorted(L):
                        if curFile.startswith(curSubject) and os.path.isfile(os.path.join(outputDirectory, 'seg', curFile)):
                            shutil.copy(os.path.join(outputDirectory, 'seg', curFile), os.path.join(outputDirectory, 'seg_bad'))

                    if quitOnError:
                        quit()    
                    else:
                        continue
                except Exception as e:
                    print("CCSeg Error: " + str(e) + ", subject " + outputBase)
                    L = os.listdir(os.path.join(outputDirectory, 'seg'))
                    CCSegUtils.mkdirSafe(os.path.join(outputDirectory, 'seg_bad'))
                    for curFile in sorted(L):
                        if curFile.startswith(curSubject) and os.path.isfile(os.path.join(outputDirectory, 'seg', curFile)):
                            shutil.copy(os.path.join(outputDirectory, 'seg', curFile), os.path.join(outputDirectory, 'seg_bad'))


                    print(traceback.print_exc(file=sys.stdout))
                    continue
                    
        if doThickness:
            if redoDone or not os.path.isfile(outputBase + "_thickness.hdf5"):
                print(curSubject + ": thickness")
                try:
                    #print outputBase + "_seg.hdf5"
                    CCSegPipeThickness.thicknessCC(outputBase, doGraphics = doGraphics, thicknessNoReg = thicknessNoReg)
                except AssertionError as e:
                    print("CCSeg thickness Error: " + str(e) + ", subject " + outputBase)
                    L = os.listdir(os.path.join(outputDirectory, 'seg'))
                    CCSegUtils.mkdirSafe(os.path.join(outputDirectory, 'seg_bad'))
                    for curFile in sorted(L):
                        if curFile.startswith(curSubject) and os.path.isfile(os.path.join(outputDirectory, 'seg', curFile)):
                            shutil.copy(os.path.join(outputDirectory, 'seg', curFile), os.path.join(outputDirectory, 'seg_bad'))
                    if quitOnError:
                        quit()
                    else:
                        continue
                except Exception as e:
                    print("CCSeg Thickness Error: " + str(e) + ", subject " + outputBase)
                    CCSegUtils.mkdirSafe(os.path.join(outputDirectory, 'seg_bad'))
                    CCSegUtils.mkdirSafe(os.path.join(outputDirectory, 'seg'))
                    L = os.listdir(os.path.join(outputDirectory, 'seg'))
                    for curFile in sorted(L):
                        if curFile.startswith(curSubject) and os.path.isfile(os.path.join(outputDirectory, 'seg', curFile)):
                            shutil.copy(os.path.join(outputDirectory, 'seg', curFile), os.path.join(outputDirectory, 'seg_bad'))

                    print(traceback.print_exc(file=sys.stdout))
                    if quitOnError:
                        quit()
                    else:
                        continue


        if doThicknessManedit and os.path.isfile(outputBase + "_seg_manedit.hdf5"):
            #print outputBase + "_seg.hdf5"
            if redoDone or not os.path.isfile(outputBase + "_thickness.hdf5"):
                print(curSubject + ": thicknessmanedit")
                try:
                    CCSegPipeThickness.thicknessCC(outputBase, doGraphics = doGraphics, thicknessNoReg = thicknessNoReg)
                except AssertionError as e:
                    print("CCSeg Error: " + str(e) + ", subject " + outputBase)
                    if quitOnError:
                        quit()
                    else:
                        continue
                except Exception as e:
                    print("CCSeg Thickness Error: " + str(e) + ", subject " + outputBase)
                    if quitOnError:
                        quit()
                    else:
                        continue
    
        if doSegToNative and not doThicknessManedit:
            print(curSubject + ": segtonative")
            #print "doing segtonative"
            if os.path.isfile(outputBase + "_seg.hdf5"):# and (redoDone or not os.path.isfile(outputBase + "_native.nii.gz")):
                try:
                    CCSegPipeSegToNative.segCCToNative(outputBase, dilateParaSagSlices = dilateParaSagSlices)
                except AssertionError as e:
                    print("CCSegToNative Error: " + str(e) + ", subject " + outputBase)
                    if quitOnError:
                        quit()
                    else:
                        continue
            #quit()    
        
        if doSegToNative and doThicknessManedit:
            #print "doing segtonative"
            if os.path.isfile(outputBase + "_seg.hdf5") and os.path.isfile(outputBase + "_seg_manedit.hdf5") and (redoDone or not os.path.isfile(outputBase + "_native.nii.gz")):
                print(curSubject + ": segtonative")
                try:
                    CCSegPipeSegToNative.segCCToNative(outputBase, dilateParaSagSlices = dilateParaSagSlices)
                except AssertionError as e:
                    print("CCSeg Error: " + str(e) + ", subject " + outputBase)
                    if quitOnError:
                        quit()
                    else:
                        continue
            #quit()    


    #quit()

    #for curSubject in inputNIFTIFiles:
#def main():


if __name__ == "__main__":
    main()