Exemple #1
0
def chain(options):
    """Create a registration chain pipeline from the given options."""

    # TODO:
    # one overall question for this entire piece of code is how
    # we are going to make sure that we can concatenate/add all
    # the transformations together. Many of the sub-registrations
    # that are performed (inter-subject registration, lsq6 using
    # multiple initial models) are applied to subsets of the entire 
    # data, making it harder to keep the mapping simple/straightforward


    chain_opts = options.chain  # type : ChainConf

    s = Stages()
    
    with open(options.chain.csv_file, 'r') as f:
        subject_info = parse_csv(rows=f, common_time_pt=options.chain.common_time_point)

    output_dir    = options.application.output_directory
    pipeline_name = options.application.pipeline_name

    pipeline_processed_dir = os.path.join(output_dir, pipeline_name + "_processed")
    pipeline_lsq12_common_dir = os.path.join(output_dir, pipeline_name + "_lsq12_" + options.chain.common_time_point_name)
    pipeline_nlin_common_dir = os.path.join(output_dir, pipeline_name + "_nlin_" + options.chain.common_time_point_name)
    pipeline_montage_dir = os.path.join(output_dir, pipeline_name + "_montage")
    
    
    pipeline_subject_info = map_over_time_pt_dict_in_Subject(
                                     lambda subj_str:  MincAtom(name=subj_str, pipeline_sub_dir=pipeline_processed_dir),
                                     subject_info)  # type: Dict[str, Subject[MincAtom]]
    
    # verify that in input files are proper MINC files, and that there 
    # are no duplicates in the filenames
    all_Minc_atoms = []  # type: List[MincAtom]
    for s_id, subj in pipeline_subject_info.items():
        for subj_time_pt, subj_filename in subj.time_pt_dict.items():
            all_Minc_atoms.append(subj_filename)
    # check_MINC_input_files takes strings, so pass along those instead of the actual MincAtoms
    check_MINC_input_files([minc_atom.path for minc_atom in all_Minc_atoms])

    if options.registration.input_space == InputSpace.lsq6 or \
        options.registration.input_space == InputSpace.lsq12:
        # the input files are not going through the lsq6 alignment. This is the place
        # where they will all be resampled using a single like file, and get the same
        # image dimensions/lengths/resolution. So in order for the subsequent stages to
        # finish (mincaverage stages for instance), all files need to have the same
        # image parameters:
        check_MINC_files_have_equal_dimensions_and_resolution([minc_atom.path for minc_atom in all_Minc_atoms],
                                                              additional_msg="Given that the input images are "
                                                                             "already in " + str(options.registration.input_space) +
                                                                             " space, all input files need to have "
                                                                             "the same dimensions/starts/step sizes.")

    if options.registration.input_space not in InputSpace.__members__.values():
        raise ValueError('unrecognized input space: %s; choices: %s' %
                         (options.registration.input_space, ','.join(InputSpace.__members__)))
    
    if options.registration.input_space == InputSpace.native:
        if options.lsq6.target_type == TargetType.bootstrap:
            raise ValueError("\nA bootstrap model is ill-defined for the registration chain. "
                             "(Which file is the 'first' input file?). Please use the --lsq6-target "
                             "flag to specify a target for the lsq6 stage, or use an initial model.")
        if options.lsq6.target_type == TargetType.pride_of_models:
            pride_of_models_dict = get_pride_of_models_mapping(pride_csv=options.lsq6.target_file,
                                                               output_dir=options.application.output_directory,
                                                               pipeline_name=options.application.pipeline_name)
            subj_id_to_subj_with_lsq6_xfm_dict = map_with_index_over_time_pt_dict_in_Subject(
                                    lambda subj_atom, time_point:
                                        s.defer(lsq6_nuc_inorm([subj_atom],
                                                               registration_targets=get_closest_model_from_pride_of_models(
                                                                                        pride_of_models_dict, time_point),
                                                               resolution=options.registration.resolution,
                                                               lsq6_options=options.lsq6,
                                                               lsq6_dir=None,  # never used since no average
                                                               # (could call this "average_dir" with None -> no avg ?)
                                                               subject_matter=options.registration.subject_matter,
                                                               create_qc_images=False,
                                                               create_average=False))[0],
                                        pipeline_subject_info)  # type: Dict[str, Subject[XfmHandler]]
        else:
            # if we are not dealing with a pride of models, we can retrieve a fixed
            # registration target for all input files:
            targets = registration_targets(lsq6_conf=options.lsq6,
                                           app_conf=options.application)
            
            # we want to store the xfm handlers in the same shape as pipeline_subject_info,
            # as such we will call lsq6_nuc_inorm for each file individually and simply extract
            # the first (and only) element from the resulting list via s.defer(...)[0].
            subj_id_to_subj_with_lsq6_xfm_dict = map_over_time_pt_dict_in_Subject(
                                         lambda subj_atom:
                                           s.defer(lsq6_nuc_inorm([subj_atom],
                                                                  registration_targets=targets,
                                                                  resolution=options.registration.resolution,
                                                                  lsq6_options=options.lsq6,
                                                                  lsq6_dir=None, # no average will be create, is just one file...
                                                                  create_qc_images=False,
                                                                  create_average=False,
                                                                  subject_matter=options.registration.subject_matter)
                                                   )[0],
                                         pipeline_subject_info)  # type: Dict[str, Subject[XfmHandler]]

        # create verification images to show the 6 parameter alignment
        montageLSQ6 = pipeline_montage_dir + "/quality_control_montage_lsq6.png"
        # TODO, base scaling factor on resolution of initial model or target
        filesToCreateImagesFrom = []
        for subj_id, subj in subj_id_to_subj_with_lsq6_xfm_dict.items():
            for time_pt, subj_time_pt_xfm in subj.time_pt_dict.items():
                filesToCreateImagesFrom.append(subj_time_pt_xfm.resampled)

        # TODO it's strange that create_quality_control_images gets the montage directory twice
        # TODO (in montages=output=montageLSQ6 and in montage_dir), suggesting a weirdness in create_q_c_images
        lsq6VerificationImages = s.defer(create_quality_control_images(filesToCreateImagesFrom,
                                                                       montage_output=montageLSQ6,
                                                                       montage_dir=pipeline_montage_dir,
                                                                       message=" the input images after the lsq6 alignment"))

    # NB currently LSQ6 expects an array of files, but we have a map.
    # possibilities:
    # - note that pairwise is enough (except for efficiency -- redundant blurring, etc.)
    #   and just use the map fn above with an LSQ6 fn taking only a single source
    # - rewrite LSQ6 to use such a (nested) map
    # - write conversion which creates a tagged array from the map, performs LSQ6,
    #   and converts back
    # - write 'over' which takes a registration, a data structure, and 'get/set' fns ...?
    

    # Intersubject registration: LSQ12/NLIN registration of common-timepoint images
    # The assumption here is that all these files are roughly aligned. Here is a toy
    # schematic of what happens. In this example, the common timepoint is set timepoint 2: 
    #
    #                            ------------
    # subject A    A_time_1   -> | A_time_2 | ->   A_time_3
    # subject B    B_time_1   -> | B_time_2 | ->   B_time_3
    # subject C    C_time_1   -> | C_time_2 | ->   C_time_3
    #                            ------------
    #                                 |
    #                            group_wise registration on time point 2
    #

    # dictionary that holds the transformations from the intersubject images
    # to the final common space average
    intersubj_img_to_xfm_to_common_avg_dict = {}  # type: Dict[MincAtom, XfmHandler]
    if options.registration.input_space in (InputSpace.lsq6, InputSpace.lsq12):
        # no registrations have been performed yet, so we can point to the input files
        s_id_to_intersubj_img_dict = { s_id : subj.intersubject_registration_image
                          for s_id, subj in pipeline_subject_info.items() }
    else:
        # lsq6 aligned images
        # When we ran the lsq6 alignment, we stored the XfmHandlers in the Subject dictionary. So when we call
        # xfmhandler.intersubject_registration_image, this returns an XfmHandler. From which
        # we want to extract the resampled file (in order to continue the registration with)
        s_id_to_intersubj_img_dict = { s_id : subj_with_xfmhandler.intersubject_registration_image.resampled
                          for s_id, subj_with_xfmhandler in subj_id_to_subj_with_lsq6_xfm_dict.items() }
    
    if options.application.verbose:
        print("\nImages that are used for the inter-subject registration:")
        print("ID\timage")
        for subject in s_id_to_intersubj_img_dict:
            print(subject + '\t' + s_id_to_intersubj_img_dict[subject].path)

    # determine what configuration to use for the non linear registration
    nonlinear_configuration = get_nonlinear_configuration_from_options(options.nlin.nlin_protocol,
                                                                       options.nlin.reg_method,
                                                                       options.registration.resolution)

    if options.registration.input_space in [InputSpace.lsq6, InputSpace.native]:
        intersubj_xfms = s.defer(lsq12_nlin_build_model(imgs=list(s_id_to_intersubj_img_dict.values()),
                                                lsq12_conf=options.lsq12,
                                                nlin_conf=nonlinear_configuration,
                                                resolution=options.registration.resolution,
                                                lsq12_dir=pipeline_lsq12_common_dir,
                                                nlin_dir=pipeline_nlin_common_dir,
                                                nlin_prefix="common"))
                                                #, like={atlas_from_init_model_at_this_tp}
    elif options.registration.input_space == InputSpace.lsq12:
        #TODO: write reader that creates a mincANTS configuration out of an input protocol
        # if we're starting with files that are already aligned with an affine transformation
        # (overall scaling is also dealt with), then the target for the non linear registration
        # should be the averge of the current input files.
        first_nlin_target = s.defer(mincaverage(imgs=list(s_id_to_intersubj_img_dict.values()),
                                                name_wo_ext="avg_of_input_files",
                                                output_dir=pipeline_nlin_common_dir))
        intersubj_xfms = s.defer(mincANTS_NLIN_build_model(imgs=list(s_id_to_intersubj_img_dict.values()),
                                                   initial_target=first_nlin_target,
                                                   nlin_dir=pipeline_nlin_common_dir,
                                                   conf=nonlinear_configuration))


    intersubj_img_to_xfm_to_common_avg_dict = { xfm.source : xfm for xfm in intersubj_xfms.output }

    # create one more convenience data structure: a mapping from subject_ID to the xfm_handler
    # that contains the transformation from the subject at the common time point to the
    # common time point average.
    subj_ID_to_xfm_handler_to_common_avg = {}
    for s_id, subj_at_common_tp in s_id_to_intersubj_img_dict.items():
        subj_ID_to_xfm_handler_to_common_avg[s_id] = intersubj_img_to_xfm_to_common_avg_dict[subj_at_common_tp]

    # create verification images to show the inter-subject  alignment
    montage_inter_subject = pipeline_montage_dir + "/quality_control_montage_inter_subject_registration.png"
    avg_and_inter_subject_images = []
    avg_and_inter_subject_images.append(intersubj_xfms.avg_img)
    for xfmh in intersubj_xfms.output:
        avg_and_inter_subject_images.append(xfmh.resampled)

    inter_subject_verification_images = s.defer(create_quality_control_images(
                                                  imgs=avg_and_inter_subject_images,
                                                  montage_output=montage_inter_subject,
                                                  montage_dir=pipeline_montage_dir,
                                                  message=" the result of the inter-subject alignment"))

    if options.application.verbose:
        print("\nTransformations for intersubject images to final nlin common space:")
        print("MincAtom\ttransformation")
        for subj_atom, xfm_handler in intersubj_img_to_xfm_to_common_avg_dict.items():
            print(subj_atom.path + '\t' + xfm_handler.xfm.path)


    ## within-subject registration
    # In the toy scenario below: 
    # subject A    A_time_1   ->   A_time_2   ->   A_time_3
    # subject B    B_time_1   ->   B_time_2   ->   B_time_3
    # subject C    C_time_1   ->   C_time_2   ->   C_time_3
    # 
    # The following registrations are run:
    # 1) A_time_1   ->   A_time_2
    # 2) A_time_2   ->   A_time_3
    #
    # 3) B_time_1   ->   B_time_2
    # 4) B_time_2   ->   B_time_3
    #
    # 5) C_time_1   ->   C_time_2
    # 6) C_time_2   ->   C_time_3    

    subj_id_to_Subjec_for_within_dict = pipeline_subject_info
    if options.registration.input_space == InputSpace.native:
        # we started with input images that were not aligned whatsoever
        # in this case we should use the images that were rigidly
        # aligned files to continue the within-subject registration with
        # # type: Dict[str, Subject[XfmHandler]]
        subj_id_to_Subjec_for_within_dict = map_over_time_pt_dict_in_Subject(lambda x: x.resampled,
                                                                             subj_id_to_subj_with_lsq6_xfm_dict)

    if options.application.verbose:
        print("\n\nWithin subject registrations:")
        for s_id, subj in subj_id_to_Subjec_for_within_dict.items():
            print("ID: ", s_id)
            for time_pt, subj_img in subj.time_pt_dict.items():
                print(time_pt, " ", subj_img.path)
            print("\n")

    # dictionary that maps subject IDs to a list containing:
    # ( [(time_pt_n, time_pt_n+1, XfmHandler_from_n_to_n+1), ..., (,,,)],
    #   index_of_common_time_pt)
    chain_xfms = { s_id : s.defer(intrasubject_registrations(
                                    subj=subj,
                                    linear_conf=default_lsq12_multilevel_minctracc,
                                    nlin_conf=mincANTS_default_conf.replace(
                                        file_resolution=options.registration.resolution,
                                        iterations="100x100x100x50")))
                   for s_id, subj in subj_id_to_Subjec_for_within_dict.items() }

    # create a montage image for each pair of time points
    for s_id, output_from_intra in chain_xfms.items():
        for time_pt_n, time_pt_n_plus_1, transform in output_from_intra[0]:
            montage_chain = pipeline_montage_dir + "/quality_control_chain_ID_" + s_id + \
                            "_timepoint_" + str(time_pt_n) + "_to_" + str(time_pt_n_plus_1) + ".png"
            chain_images = [transform.resampled, transform.target]
            chain_verification_images = s.defer(create_quality_control_images(chain_images,
                                                                              montage_output=montage_chain,
                                                                              montage_dir=pipeline_montage_dir,
                                                                              message="the alignment between ID " + s_id + " time point " +
                                                                                      str(time_pt_n) + " and " + str(time_pt_n_plus_1)))

    if options.application.verbose:
        print("\n\nTransformations gotten from the intrasubject registrations:")
        for s_id, output_from_intra in chain_xfms.items():
            print("ID: ", s_id)
            for time_pt_n, time_pt_n_plus_1, transform in output_from_intra[0]:
                print("Time point: ", time_pt_n, " to ", time_pt_n_plus_1, " trans: ", transform.xfm.path)
            print("\n")

    ## stats
    #
    # The statistic files we want to create are the following:
    # 1) subject <----- subject_common_time_point                              (resampled to common average)
    # 2) subject <----- subject_common_time_point <- common_time_point_average (incorporates inter subject differences)
    # 3) subject_time_point_n <----- subject_time_point_n+1                    (resampled to common average)

    # create transformation from each subject to the final common time point average,
    # and from each subject to the subject's common time point
    (non_rigid_xfms_to_common_avg, non_rigid_xfms_to_common_subj) = s.defer(get_chain_transforms_for_stats(subj_id_to_Subjec_for_within_dict,
                                                                            intersubj_img_to_xfm_to_common_avg_dict,
                                                                            chain_xfms))

    # Ad 1) provide transformations from the subject's common time point to each subject
    #       These are temporary, because they still need to be resampled into the
    #       average common time point space
    determinants_from_subject_common_to_subject = map_over_time_pt_dict_in_Subject(
        lambda xfm: s.defer(determinants_at_fwhms(xfms=[s.defer(invert_xfmhandler(xfm))],
                                                  inv_xfms=[xfm],  # determinants_at_fwhms now vectorized-unhelpful here
                                                  blur_fwhms=options.stats.stats_kernels)),
        non_rigid_xfms_to_common_subj)
    # the content of determinants_from_subject_common_to_subject is:
    #
    # {subject_ID : Subject(inter_subject_time_pt, time_pt_dict)
    #
    # where time_pt_dict contains:
    #
    # {time_point : Tuple(List[Tuple(float, Tuple(MincAtom, MincAtom))],
    #                     List[Tuple(float, Tuple(MincAtom, MincAtom))])
    #
    # And to be a bit more verbose:
    #
    # {time_point : Tuple(relative_stat_files,
    #                     absolute_stat_files)
    #
    # where either the relative_stat_files or the absolute_stat_files look like:
    #
    # [blur_kernel_1, (determinant_file_1, log_of_determinant_file_1),
    #  ...,
    #  blur_kernel_n, (determinant_file_n, log_of_determinant_file_n)]
    #
    # Now the only thing we really want to do, is to resample the actual log
    # determinants that were generated into the space of the common average.
    # To make that a little easier, I'll create a mapping that will contain:
    #
    # {subject_ID: Subject(intersubject_timepoint, {time_pt_1: [stat_file_1, ..., stat_file_n],
    #                                               ...,
    #                                               time_pt_n: [stat_file_1, ..., stat_file_n]}
    # }
    for s_id, subject_with_determinants in determinants_from_subject_common_to_subject.items():
        transform_from_common_subj_to_common_avg = subj_ID_to_xfm_handler_to_common_avg[s_id].xfm
        for time_pt, determinant_info in subject_with_determinants.time_pt_dict.items():
            # here, each determinant_info is a DataFrame where each row contains
            # 'abs_det', 'nlin_det', 'log_nlin_det', 'log_abs_det', 'fwhm' fields
            # of the log-determinants, blurred at various fwhms (corresponding to different rows)
            for _ix, row in determinant_info.iterrows():
                for log_det_file_to_resample in (row.log_full_det, row.log_nlin_det):
                    # TODO the MincAtoms corresponding to the resampled files are never returned
                    new_name_wo_ext = log_det_file_to_resample.filename_wo_ext + "_resampled_to_common"
                    s.defer(mincresample(img=log_det_file_to_resample,
                                         xfm=transform_from_common_subj_to_common_avg,
                                         like=log_det_file_to_resample,
                                         new_name_wo_ext=new_name_wo_ext,
                                         subdir="stats-volumes"))

    # Ad 2) provide transformations from the common avg to each subject. That's the
    #       inverse of what was provided by get_chain_transforms_for_stats()
    determinants_from_common_avg_to_subject = map_over_time_pt_dict_in_Subject(
        lambda xfm: s.defer(determinants_at_fwhms(xfms=[s.defer(invert_xfmhandler(xfm))],
                                                  inv_xfms=[xfm],  # determinants_at_fwhms now vectorized-unhelpful here
                                                  blur_fwhms=options.stats.stats_kernels)),
        non_rigid_xfms_to_common_avg)

    # TODO don't just return an (unnamed-)tuple here
    return Result(stages=s, output=Namespace(non_rigid_xfms_to_common=non_rigid_xfms_to_common_avg,
                                             determinants_from_common_avg_to_subject=determinants_from_common_avg_to_subject,
                                             determinants_from_subject_common_to_subject=determinants_from_subject_common_to_subject))
Exemple #2
0
def get_chain_transforms_for_stats(pipeline_subject_info, intersubj_xfms_dict, chain_xfms_dict):
    """
    pipeline_subject_info

    intersubj_xfms_dict

    chain_xfms_dict -- {subject_ID : ( [ (time_point_n, time_point_n+1, XfmHandler(time_point -> time_point + 1)),
                                         ..., (,,,) ],
                                       index_of_common_time_point) }


    This function takes a subject mapping (with timepoints to MincAtoms) and returns a
    subject mapping of timepoints to `XfmHandler`s. Those transformations for
    each subject will contain the non-rigid transformation to the common time point average
    
    chain_xfms_dict maps subject_ids to a tuple containing a list of tuples
    (time_point_n, time_point_n+1, transformation)
    and the index to the common time point in that list
    """

    s = Stages()
    dict_transforms_to_common_avg = {}
    dict_transforms_to_subject_common_tp = {}
    dict_transforms_from_common_tp_to_common_avg = {}
    for s_id, subj in pipeline_subject_info.items():
        # dictionary: {time_pt : XfmHandler_time_pt_to_final_common_avg}
        trans_to_final_common_avg_dict = {}
        # dictionary: {time_pt : XfmHandler_time_pt_to_subject_common_time_pt}
        trans_to_subject_common_time_pt = {}



        ##############################################################
        #                -------------
        #  time_1   ... | time_common | ...   time_n
        #                -------------
        #
        # the transformation for the common time point is easy,
        # intersubj_xfms_dict[subj.intersubject_registration_image]
        # returns the XfmHandler from the subject common time point
        # to the common time point average
        trans_to_final_common_avg_dict[subj.intersubject_registration_time_pt] = \
            intersubj_xfms_dict[subj.intersubject_registration_image]
        # there is no transform from the common time point to the common time
        # point. Technically it is the identity transformation, but there is
        # no use in generating a stats file from the identity transformation,
        # so we simply won't generate anything

        chain_transforms, index_of_common_time_pt = chain_xfms_dict[s_id]

        # will hold the XfmHandler from current to average of common time pt
        current_xfm_to_common_avg = intersubj_xfms_dict[subj.intersubject_registration_image]
        # and the XfmHandler from current to the subject common time point
        # starts out as None (technically the identity transformation)
        current_xfm_to_common_subject = None

        # we start at the common time point and are going forward at this point
        # so we will assign the concatenated transform to the target of each 
        # transform we are adding (which is why we take the inverse)
        #
        #                       - - - - - - - - >
        #  time_1   ...   time_common   ...   time_n
        #
        #
        for time_pt_n, time_pt_n_plus_1, transform in chain_transforms[index_of_common_time_pt:]:
            current_xfm_to_common_avg = s.defer(concat_xfmhandlers([s.defer(invert_xfmhandler(transform)), current_xfm_to_common_avg],
                                                               name="id_%s_pt_%s_to_common_avg" % (s_id, time_pt_n_plus_1)))
            # we are moving away from the common time point. That means that the transformation
            # we are adding here is the inverse of n -> n+1, and should be added to time point n+1
            trans_to_final_common_avg_dict[time_pt_n_plus_1] = current_xfm_to_common_avg

            if current_xfm_to_common_subject == None:
                current_xfm_to_common_subject = s.defer(invert_xfmhandler(transform))
            else:
                current_xfm_to_common_subject = s.defer(concat_xfmhandlers([s.defer(invert_xfmhandler(transform)), current_xfm_to_common_subject],
                                                               name="id_%s_pt_%s_to_common_subject" % (s_id, time_pt_n_plus_1)))
            trans_to_subject_common_time_pt[time_pt_n_plus_1] = current_xfm_to_common_subject

        # we need to do something similar moving backwards: make sure to reset
        # the current_xfm_to_common_avg here!
        #
        #    < - - - - - - - -
        #  time_1   ...   time_common   ...   time_n
        #
        #
        current_xfm_to_common_avg = intersubj_xfms_dict[subj.intersubject_registration_image]
        current_xfm_to_common_subject = None
        # we have to be careful here... if the index_of_common_time_pt is 0 (i.e. all images are
        # registered towards the first file in the time line, the following command will call:
        # .... chain_transforms[-1::-1] and that in turn will start at the end of the list
        # because -1 wraps around. To prevent this case, we ensure that the index_of_common_time_pt
        # is greater than 0
        if index_of_common_time_pt > 0:
            for time_pt_n, time_pt_n_plus_1, transform in chain_transforms[index_of_common_time_pt-1::-1]:
                current_xfm_to_common_avg = s.defer(concat_xfmhandlers([transform, current_xfm_to_common_avg],
                                                                   name="id_%s_pt_%s_to_common_avg" % (s_id, time_pt_n)))
                trans_to_final_common_avg_dict[time_pt_n] = current_xfm_to_common_avg

                if current_xfm_to_common_subject == None:
                    current_xfm_to_common_subject = transform
                else:
                    current_xfm_to_common_subject = s.defer(concat_xfmhandlers([transform, current_xfm_to_common_subject],
                                                                   name="id_%s_pt_%s_to_common_subject" % (s_id, time_pt_n)))
                trans_to_subject_common_time_pt[time_pt_n] = current_xfm_to_common_subject

        new_subj_to_common_avg = Subject(intersubject_registration_time_pt = subj.intersubject_registration_time_pt,
                           time_pt_dict = trans_to_final_common_avg_dict)
        dict_transforms_to_common_avg[s_id] = new_subj_to_common_avg

        new_subj_to_common_subj = Subject(intersubject_registration_time_pt = subj.intersubject_registration_time_pt,
                          time_pt_dict = trans_to_subject_common_time_pt)
        dict_transforms_to_subject_common_tp[s_id] = new_subj_to_common_subj
    return Result(stages=s, output=(dict_transforms_to_common_avg, dict_transforms_to_subject_common_tp))