registration_interface.py

import argparse
import os
import numpy as np
import time
import shutil
import registration as reg
from dagreg import dagRegistration
from tradreg import *
from boldli import ImageManipulatingLibrary as mil

def main():
    # Set up argument parser
    parser = argparse.ArgumentParser(description="Perform motion correction on time-series images.")
    # image filenames
    parser.add_argument('-i', '--inputFn', type=str, help='Full path to the name of the file to correct')
    parser.add_argument('-o', '--outputFn', type=str, help='The name of the file to save the correction to (within the directory containing the original file.')
    # which type of global volume registration framework
    parser.add_argument('-t', '--correctionType', type=str, help='Specify which type of correction to run. '
                        +'Options include: traditional or dag')
    # which type of registration
    parser.add_argument('-r', '--registrationType', type=str, help='Specify which type of registration to use: affine or nonlinear')

    # now parse the arguments
    args = parser.parse_args()
    print(args)

    # image filename
    origFn = args.inputFn.replace("//", '/')
    baseDir = origFn.rsplit("/", 1)[0]+'/'
    print(origFn)
    print(baseDir)

    # Make the directory for the transform parameters
    transformDir = os.path.join(baseDir,'transforms/')
    if not os.path.exists(transformDir):
        os.mkdir(transformDir)

    # Make the directory for the original image volunmes
    originalDir = os.path.join(baseDir, 'original/')
    if not os.path.exists(originalDir):
        os.mkdir(originalDir)

    # Make the directory for the registered images
    if args.correctionType == 'traditional' or args.correctionType == 'dag':
        registeredDir = os.path.join(baseDir, args.correctionType+"/")
        if not os.path.exists(registeredDir):
            os.mkdir(registeredDir)

    # divide the image into timepoints
    timepointFns = reg.expandTimepoints(origFn, originalDir)

    # Select the specified motion correction algorithm
    registeredFns = []
    
    if args.correctionType == 'traditional':
        """
        # Traditional volume registration: align each image to the 
        # first timepoint
        """
        # register the images sequentially
        templateImg = timepointFns[0]
        registeredFns = volumeRegistration(templateImg, timepointFns, registeredDir, transformDir, regType=args.registrationType)
        registeredFns.append(outputDir+'000.nii.gz')
        registeredFns = sorted(registeredFns)

    elif args.correctionType == 'dag':
        """
        #DAG-based registration: Treat the image series as a DAG.
        # - Treat the first volume as the template volume
        # - Register the second volume to the first
        # - Use the transformation generated by this registration to
        #   initialize the registration between the third volume and
        #   the first volume
        # - Repeat the previous step for the remaining volumes in the
        #   image sequence
        """
        # register the images using dag correction
        registeredFns = dagRegistration(timepointFns, registeredDir, transformDir, transformType=args.registrationType)

    else:
        print("Error: the type of motion correction entered is not currently supported.")
        print("       Entered:", args.correctionType)
    

    registeredFns = ["./dag/"+str(i).zfill(3)+".nii.gz" for i in range(10)]

    # load the template image
    img, coord = mil.loadBOLD(origFn)

    print(registeredFns)
    print(baseDir)

    # combine the registered timepoints into 1 file
    comboFn = os.path.join(baseDir,args.outputFn)
    reg.stackNiftis(registeredFns, coord, comboFn)

    return origFn, args.correctionType


if __name__ == "__main__":
    # very crude numpy version check
    npVer = np.__version__
    npVerList = [int(i) for i in npVer.split('.')]
    if npVerList[1] < 12:
        sys.exit("Warning: the version for numpy is "+np.__version__+".\nPlease update to at least version 1.12.1 to use this pipeline.")

    # Want to time how long it takes to register the image        
    startTime = time.time()
    subjFn, method = main()
    endTime = time.time() - startTime

    # Convert the amount of time taken from a float to a string
    totalDays = np.floor(endTime/24.0/60.0/60.0)
    endTime = endTime%(24.0*60.0*60.0)
    totalHours = np.floor(endTime/60.0/60.0)
    endTime = endTime%(60.0*60.0)
    totalMins = np.floor(endTime/60.0)
    totalSecs = endTime%60.0
    # Notify the user of the runtime
    print("Total run time:",totalDays,"days,",totalHours,"hours,",totalMins,"minutes,",totalSecs)

    # write the time to a file
    subj = subjFn.split("/")[-2]
    baseDir = subjFn.split(subj)[0]
    fn = baseDir+"timeToRun.csv"
    timeLine = subj+", "+ method+", "+ str(totalDays).zfill(2) + ":" + str(totalHours).zfill(2) + ":" + str(totalMins).zfill(2) + ":" + str(totalSecs).zfill(2) + "\n"
    # if the file aready exists
    if not os.path.isfile(fn):
        # open the file in write
        with open(fn, "w") as file:
            headerLine = "Subject, Method, RunTime (DD:HH:MM:SS)\n"
            # write the header line
            file.write(headerLine)
            # write the time line
            file.write(timeLine)

    else:
        with open(fn, "a+") as file:
            # write the time line
            file.write(timeLine)