def _run(self):

        ph.print_title("Multi-Component Reconstruction")

        self._reconstructions = [None] * len(self._stacks)

        for i in range(0, len(self._stacks)):
            ph.print_subtitle("Multi-Component Reconstruction -- "
                              "Stack %d/%d" % (i + 1, len(self._stacks)))
            stack = self._stacks[i]
            self._reconstruction_method.set_stacks([stack])
            self._reconstruction_method.run()
            self._reconstructions[i] = st.Stack.from_stack(
                self._reconstruction_method.get_reconstruction())
            self._reconstructions[i].set_filename(stack.get_filename() +
                                                  self._suffix)
Beispiel #2
0
    def print_arguments(self, args, title="Configuration:"):
        ph.print_title(title)
        for arg in sorted(vars(args)):
            ph.print_info("%s: " % (arg), newline=False)
            vals = getattr(args, arg)

            if type(vals) is list:
                # print list element in new lines, unless only one entry in list
                # if len(vals) == 1:
                #     print(vals[0])
                # else:
                print("")
                for val in vals:
                    print("\t%s" % val)
            else:
                print(vals)
        print("\nNiftyMIC version: %s" % niftymic.__version__)
        ph.print_line_separator(add_newline=False)
        print("")
    def _run(self):

        ph.print_title("Volume-to-Volume Registration")

        self._registration_method.set_moving(self._reference)

        for i in range(0, len(self._stacks)):
            txt = "Volume-to-Volume Registration -- " \
                "Stack %d/%d" % (i + 1, len(self._stacks))
            if self._verbose:
                ph.print_subtitle(txt)
            else:
                ph.print_info(txt)

            self._registration_method.set_fixed(self._stacks[i])
            self._registration_method.run()

            # Update position of stack
            transform_sitk = \
                self._registration_method.\
                get_registration_transform_sitk()
            self._stacks[i].update_motion_correction(transform_sitk)
Beispiel #4
0
    def _run(self, debug=0):
        ph.print_title(
            "Hierarchical SliceSet2V-Registration")

        N_stacks = len(self._stacks)

        self._registration_method.set_moving(self._reference)

        for i_stack, stack in enumerate(self._stacks):
            n_slices = stack.get_number_of_slices()
            for i in range(self._interleave):
                package = list(np.arange(i, n_slices, self._interleave))
                if len(package) / 2 >= self._min_slices:
                    indices_splits = self._recursive_split(
                        package, [], self._min_slices)
                else:
                    indices_splits = [package]

                prefix = "Hierarchical S2V-Reg: " \
                    "Stack %d/%d (%s) -- Interleave %d/%d --" % (
                        i_stack + 1, len(self._stacks), stack.get_filename(),
                        i + 1, self._interleave,
                    )
                if debug:
                    ph.print_subtitle(
                        "%s %d splits: %s" % (
                            prefix, len(indices_splits), indices_splits),
                    )

                ss2vreg = SliceSetToVolumeRegistration(
                    print_prefix=prefix,
                    stack=stack,
                    reference=self._reference,
                    registration_method=self._registration_method,
                    slice_set_indices=indices_splits,
                    verbose=self._verbose,
                )
                ss2vreg.run()
Beispiel #5
0
    def _run(self):

        ph.print_title("Two-step S2V-Registration and SRR Reconstruction")

        s2vreg = SliceToVolumeRegistration(
            stacks=self._stacks,
            reference=self._reference,
            registration_method=self._registration_method,
            verbose=False,
            interleave=self._interleave,
        )

        reference = self._reference

        for cycle in range(0, self._cycles):

            if cycle == 0 and self._use_hierarchical_registration:
                hs2vreg = HieararchicalSliceSetRegistration(
                    stacks=self._stacks,
                    reference=reference,
                    registration_method=self._registration_method,
                    interleave=self._interleave,
                    viewer=self._viewer,
                    min_slices=1,
                    verbose=False,
                )
                hs2vreg.run()
                self._computational_time_registration += \
                    hs2vreg.get_computational_time()
            else:
                # Slice-to-volume registration step
                s2vreg.set_reference(reference)
                s2vreg.set_print_prefix("Cycle %d/%d: " %
                                        (cycle + 1, self._cycles))
                s2vreg.run()

            self._computational_time_registration += \
                s2vreg.get_computational_time()

            # Reject misregistered slices
            if self._outlier_rejection:
                ph.print_subtitle("Slice Outlier Rejection (%s < %g)" % (
                    self._threshold_measure, self._thresholds[cycle]))
                outlier_rejector = outre.OutlierRejector(
                    stacks=self._stacks,
                    reference=self._reference,
                    threshold=self._thresholds[cycle],
                    measure=self._threshold_measure,
                    verbose=True,
                )
                outlier_rejector.run()
                self._reconstruction_method.set_stacks(
                    outlier_rejector.get_stacks())

                if len(self._stacks) == 0:
                    raise RuntimeError(
                        "All slices of all stacks were rejected "
                        "as outliers. Volumetric reconstruction is aborted.")

            # SRR step
            if cycle < self._cycles - 1:
                # ---------------- Perform Image Reconstruction ---------------
                ph.print_subtitle("Volumetric Image Reconstruction")
                if isinstance(
                    self._reconstruction_method,
                    sda.ScatteredDataApproximation
                ):
                    self._reconstruction_method.set_sigma(self._alphas[cycle])
                else:
                    self._reconstruction_method.set_alpha(self._alphas[cycle])
                self._reconstruction_method.run()

                self._computational_time_reconstruction += \
                    self._reconstruction_method.get_computational_time()

                reference = self._reconstruction_method.get_reconstruction()

                # ------------------ Perform Image Mask SDA -------------------
                ph.print_subtitle("Volumetric Image Mask Reconstruction")
                SDA = sda.ScatteredDataApproximation(
                    self._stacks,
                    reference,
                    sigma=self._sigma_sda_mask,
                    sda_mask=True,
                )
                SDA.run()

                # reference contains updated mask based on SDA
                reference = SDA.get_reconstruction()

                # -------------------- Store Reconstruction -------------------
                filename = "Iter%d_%s" % (
                    cycle + 1,
                    self._reconstruction_method.get_setting_specific_filename()
                )
                self._reconstructions.insert(0, st.Stack.from_stack(
                    reference, filename=filename))

                if self._verbose:
                    sitkh.show_stacks(self._reconstructions,
                                      segmentation=self._reference,
                                      viewer=self._viewer)
Beispiel #6
0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Register an obtained reconstruction (moving) "
        "to a template image/space (fixed) using rigid registration. "
        "The resulting registration can optionally be applied to previously "
        "obtained motion correction slice transforms so that a volumetric "
        "reconstruction is possible in the (standard anatomical) space "
        "defined by the fixed.", )
    input_parser.add_fixed(required=True)
    input_parser.add_moving(
        required=True,
        nargs="+",
        help="Specify moving image to be warped to fixed space. "
        "If multiple images are provided, all images will be transformed "
        "uniformly according to the registration obtained for the first one.")
    input_parser.add_dir_output(required=True)
    input_parser.add_dir_input()
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_search_angle(default=180)
    input_parser.add_option(
        option_string="--transform-only",
        type=int,
        help="Turn on/off functionality to transform moving image(s) to fixed "
        "image only, i.e. no resampling to fixed image space",
        default=0)
    input_parser.add_option(
        option_string="--write-transform",
        type=int,
        help="Turn on/off functionality to write registration transform",
        default=0)
    input_parser.add_verbose(default=0)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    use_reg_aladin_for_refinement = True

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")
    data_reader = dr.MultipleImagesReader(args.moving, suffix_mask="_mask")
    data_reader.read_data()
    moving = data_reader.get_data()

    data_reader = dr.MultipleImagesReader([args.fixed], suffix_mask="_mask")
    data_reader.read_data()
    fixed = data_reader.get_data()[0]

    # -------------------Register Reconstruction to Template-------------------
    ph.print_title("Register Reconstruction to Template")

    # Define search angle ranges for FLIRT in all three dimensions
    search_angles = [
        "-searchr%s -%d %d" % (x, args.search_angle, args.search_angle)
        for x in ["x", "y", "z"]
    ]
    search_angles = (" ").join(search_angles)
    options_args = []
    options_args.append(search_angles)
    # cost = "mutualinfo"
    # options_args.append("-searchcost %s -cost %s" % (cost, cost))
    registration = regflirt.FLIRT(
        fixed=moving[0],
        moving=fixed,
        # use_fixed_mask=True,
        # use_moving_mask=True,  # moving mask only seems to work for SB cases
        registration_type="Rigid",
        use_verbose=False,
        options=(" ").join(options_args),
    )
    ph.print_info("Run Registration (FLIRT) ... ", newline=False)
    registration.run()
    print("done")
    transform_sitk = registration.get_registration_transform_sitk()

    if args.write_transform:
        path_to_transform = os.path.join(args.dir_output,
                                         "registration_transform_sitk.txt")
        sitk.WriteTransform(transform_sitk, path_to_transform)

    # Apply rigidly transform to align reconstruction (moving) with template
    # (fixed)
    for m in moving:
        m.update_motion_correction(transform_sitk)

        # Additionally, use RegAladin for more accurate alignment
        # Rationale: FLIRT has better capture range, but RegAladin seems to
        # find better alignment once it is within its capture range.
        if use_reg_aladin_for_refinement:
            registration = niftyreg.RegAladin(
                fixed=m,
                use_fixed_mask=True,
                moving=fixed,
                registration_type="Rigid",
                use_verbose=False,
            )
            ph.print_info("Run Registration (RegAladin) ... ", newline=False)
            registration.run()
            print("done")
            transform2_sitk = registration.get_registration_transform_sitk()
            m.update_motion_correction(transform2_sitk)
            transform_sitk = sitkh.get_composite_sitk_affine_transform(
                transform2_sitk, transform_sitk)

    if args.transform_only:
        for m in moving:
            m.write(args.dir_output, write_mask=False)
        ph.exit()

    # Resample reconstruction (moving) to template space (fixed)
    warped_moving = [
        m.get_resampled_stack(fixed.sitk, interpolator="Linear")
        for m in moving
    ]

    for wm in warped_moving:
        wm.set_filename(wm.get_filename() + "ResamplingToTemplateSpace")

        if args.verbose:
            sitkh.show_stacks([fixed, wm], segmentation=fixed)

        # Write resampled reconstruction (moving)
        wm.write(args.dir_output, write_mask=False)

    if args.dir_input is not None:
        data_reader = dr.ImageSlicesDirectoryReader(
            path_to_directory=args.dir_input, suffix_mask=args.suffix_mask)
        data_reader.read_data()
        stacks = data_reader.get_data()

        for i, stack in enumerate(stacks):
            stack.update_motion_correction(transform_sitk)
            ph.print_info("Stack %d/%d: All slice transforms updated" %
                          (i + 1, len(stacks)))

            # Write transformed slices
            stack.write(
                os.path.join(args.dir_output, "motion_correction"),
                write_mask=True,
                write_slices=True,
                write_transforms=True,
                suffix_mask=args.suffix_mask,
            )

    elapsed_time_total = ph.stop_timing(time_start)

    # Summary
    ph.print_title("Summary")
    print("Computational Time: %s" % (elapsed_time_total))

    return 0
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    # Read input
    input_parser = InputArgparser(
        description="Volumetric MRI reconstruction framework to reconstruct "
        "an isotropic, high-resolution 3D volume from multiple "
        "motion-corrected (or static) stacks of low-resolution slices.", )
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks()
    input_parser.add_dir_input_mc()
    input_parser.add_output(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_target_stack(default=None)
    input_parser.add_extra_frame_target(default=10)
    input_parser.add_isotropic_resolution(default=None)
    input_parser.add_intensity_correction(default=1)
    input_parser.add_reconstruction_space(default=None)
    input_parser.add_minimizer(default="lsmr")
    input_parser.add_iter_max(default=10)
    input_parser.add_reconstruction_type(default="TK1L2")
    input_parser.add_data_loss(default="linear")
    input_parser.add_data_loss_scale(default=1)
    input_parser.add_alpha(default=0.01  # TK1L2
                           # default=0.006  #TVL2, HuberL2
                           )
    input_parser.add_rho(default=0.5)
    input_parser.add_tv_solver(default="PD")
    input_parser.add_pd_alg_type(default="ALG2")
    input_parser.add_iterations(default=15)
    input_parser.add_log_config(default=1)
    input_parser.add_use_masks_srr(default=0)
    input_parser.add_slice_thicknesses(default=None)
    input_parser.add_verbose(default=0)
    input_parser.add_viewer(default="itksnap")
    input_parser.add_argument(
        "--mask",
        "-mask",
        action='store_true',
        help="If given, input images are interpreted as image masks. "
        "Obtained volumetric reconstruction will be exported in uint8 format.")
    input_parser.add_argument(
        "--sda",
        "-sda",
        action='store_true',
        help="If given, the volume is reconstructed using "
        "Scattered Data Approximation (Vercauteren et al., 2006). "
        "--alpha is considered the value for the standard deviation then. "
        "Recommended value is, e.g., --alpha 0.8")

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.reconstruction_type not in ["TK1L2", "TVL2", "HuberL2"]:
        raise IOError("Reconstruction type unknown")

    if np.alltrue([not args.output.endswith(t) for t in ALLOWED_EXTENSIONS]):
        raise ValueError("output filename '%s' invalid; "
                         "allowed image extensions are: %s" %
                         (args.output, ", ".join(ALLOWED_EXTENSIONS)))

    dir_output = os.path.dirname(args.output)
    ph.create_directory(dir_output)

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    if args.verbose:
        show_niftis = []
        # show_niftis = [f for f in args.filenames]

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    if args.mask:
        filenames_masks = args.filenames
    else:
        filenames_masks = args.filenames_masks

    data_reader = dr.MultipleImagesReader(
        file_paths=args.filenames,
        file_paths_masks=filenames_masks,
        suffix_mask=args.suffix_mask,
        dir_motion_correction=args.dir_input_mc,
        stacks_slice_thicknesses=args.slice_thicknesses,
    )
    data_reader.read_data()
    stacks = data_reader.get_data()

    ph.print_info("%d input stacks read for further processing" % len(stacks))

    # Specify target stack for intensity correction and reconstruction space
    if args.target_stack is None:
        target_stack_index = 0
    else:
        filenames = ["%s.nii.gz" % s.get_filename() for s in stacks]
        filename_target_stack = os.path.basename(args.target_stack)
        try:
            target_stack_index = filenames.index(filename_target_stack)
        except ValueError as e:
            raise ValueError(
                "--target-stack must correspond to an image as provided by "
                "--filenames")

    # ---------------------------Intensity Correction--------------------------
    if args.intensity_correction and not args.mask:
        ph.print_title("Intensity Correction")
        intensity_corrector = ic.IntensityCorrection()
        intensity_corrector.use_individual_slice_correction(False)
        intensity_corrector.use_stack_mask(True)
        intensity_corrector.use_reference_mask(True)
        intensity_corrector.use_verbose(False)

        for i, stack in enumerate(stacks):
            if i == target_stack_index:
                ph.print_info("Stack %d (%s): Reference image. Skipped." %
                              (i + 1, stack.get_filename()))
                continue
            else:
                ph.print_info("Stack %d (%s): Intensity Correction ... " %
                              (i + 1, stack.get_filename()),
                              newline=False)
            intensity_corrector.set_stack(stack)
            intensity_corrector.set_reference(
                stacks[target_stack_index].get_resampled_stack(
                    resampling_grid=stack.sitk,
                    interpolator="NearestNeighbor",
                ))
            intensity_corrector.run_linear_intensity_correction()
            stacks[i] = intensity_corrector.get_intensity_corrected_stack()
            print("done (c1 = %g) " %
                  intensity_corrector.get_intensity_correction_coefficients())

    # -------------------------Volumetric Reconstruction-----------------------
    ph.print_title("Volumetric Reconstruction")

    # Reconstruction space is given isotropically resampled target stack
    if args.reconstruction_space is None:
        recon0 = stacks[target_stack_index].get_isotropically_resampled_stack(
            resolution=args.isotropic_resolution,
            extra_frame=args.extra_frame_target)
        recon0 = recon0.get_cropped_stack_based_on_mask(
            boundary_i=args.extra_frame_target,
            boundary_j=args.extra_frame_target,
            boundary_k=args.extra_frame_target,
            unit="mm",
        )

    # Reconstruction space was provided by user
    else:
        recon0 = st.Stack.from_filename(args.reconstruction_space,
                                        extract_slices=False)

        # Change resolution for isotropic resolution if provided by user
        if args.isotropic_resolution is not None:
            recon0 = recon0.get_isotropically_resampled_stack(
                args.isotropic_resolution)

        # Use image information of selected target stack as recon0 serves
        # as initial value for reconstruction
        recon0 = stacks[target_stack_index].get_resampled_stack(recon0.sitk)
        recon0 = recon0.get_stack_multiplied_with_mask()

    ph.print_info("Reconstruction space defined with %s mm3 resolution" %
                  " x ".join(["%.2f" % s for s in recon0.sitk.GetSpacing()]))

    if args.sda:
        ph.print_title("Compute SDA reconstruction")
        SDA = sda.ScatteredDataApproximation(stacks,
                                             recon0,
                                             sigma=args.alpha,
                                             sda_mask=args.mask)
        SDA.run()
        recon = SDA.get_reconstruction()
        if args.mask:
            dw.DataWriter.write_mask(recon.sitk_mask, args.output)
        else:
            dw.DataWriter.write_image(recon.sitk, args.output)

        if args.verbose:
            show_niftis.insert(0, args.output)

    else:
        if args.reconstruction_type in ["TVL2", "HuberL2"]:
            ph.print_title("Compute Initial value for %s" %
                           args.reconstruction_type)
            SRR0 = tk.TikhonovSolver(
                stacks=stacks,
                reconstruction=recon0,
                alpha=args.alpha,
                iter_max=np.min([5, args.iter_max]),
                reg_type="TK1",
                minimizer="lsmr",
                data_loss="linear",
                use_masks=args.use_masks_srr,
                # verbose=args.verbose,
            )
        else:
            ph.print_title("Compute %s reconstruction" %
                           args.reconstruction_type)
            SRR0 = tk.TikhonovSolver(
                stacks=stacks,
                reconstruction=recon0,
                alpha=args.alpha,
                iter_max=args.iter_max,
                reg_type="TK1",
                minimizer=args.minimizer,
                data_loss=args.data_loss,
                data_loss_scale=args.data_loss_scale,
                use_masks=args.use_masks_srr,
                # verbose=args.verbose,
            )
        SRR0.run()

        recon = SRR0.get_reconstruction()

        if args.reconstruction_type in ["TVL2", "HuberL2"]:
            output = ph.append_to_filename(args.output, "_initTK1L2")
        else:
            output = args.output

        if args.mask:
            mask_estimator = bm.BinaryMaskFromMaskSRREstimator(recon.sitk)
            mask_estimator.run()
            mask_sitk = mask_estimator.get_mask_sitk()
            dw.DataWriter.write_mask(mask_sitk, output)
        else:
            dw.DataWriter.write_image(recon.sitk, output)

        if args.verbose:
            show_niftis.insert(0, output)

        if args.reconstruction_type in ["TVL2", "HuberL2"]:
            ph.print_title("Compute %s reconstruction" %
                           args.reconstruction_type)
            if args.tv_solver == "ADMM":
                SRR = admm.ADMMSolver(
                    stacks=stacks,
                    reconstruction=st.Stack.from_stack(
                        SRR0.get_reconstruction()),
                    minimizer=args.minimizer,
                    alpha=args.alpha,
                    iter_max=args.iter_max,
                    rho=args.rho,
                    data_loss=args.data_loss,
                    iterations=args.iterations,
                    use_masks=args.use_masks_srr,
                    verbose=args.verbose,
                )

            else:
                SRR = pd.PrimalDualSolver(
                    stacks=stacks,
                    reconstruction=st.Stack.from_stack(
                        SRR0.get_reconstruction()),
                    minimizer=args.minimizer,
                    alpha=args.alpha,
                    iter_max=args.iter_max,
                    iterations=args.iterations,
                    alg_type=args.pd_alg_type,
                    reg_type="TV"
                    if args.reconstruction_type == "TVL2" else "huber",
                    data_loss=args.data_loss,
                    use_masks=args.use_masks_srr,
                    verbose=args.verbose,
                )
            SRR.run()
            recon = SRR.get_reconstruction()

            if args.mask:
                mask_estimator = bm.BinaryMaskFromMaskSRREstimator(recon.sitk)
                mask_estimator.run()
                mask_sitk = mask_estimator.get_mask_sitk()
                dw.DataWriter.write_mask(mask_sitk, args.output)

            else:
                dw.DataWriter.write_image(recon.sitk, args.output)

            if args.verbose:
                show_niftis.insert(0, args.output)

    if args.verbose:
        ph.show_niftis(show_niftis, viewer=args.viewer)

    ph.print_line_separator()

    elapsed_time = ph.stop_timing(time_start)
    ph.print_title("Summary")
    print("Computational Time for Volumetric Reconstruction: %s" %
          (elapsed_time))

    return 0
    def _run(self, debug=1):
        ph.print_title(
            "Hierarchical SliceSet2V-Registration and SRR Reconstruction")

        N_stacks = len(self._stacks)

        # Minimum number of stacks at which no further splitting performed
        N_min = 1
        slice_sets_indices = [None] * N_stacks
        for i, stack in enumerate(self._stacks):
            slice_sets_indices[i] = \
                self._get_slice_set_indices_per_cycle(stack, N_min=N_min)

        # Debug
        if debug:
            for i, stack in enumerate(self._stacks):
                print("Stack %d/%d:" % (i + 1, N_stacks))
                for k, v in six.iteritems(slice_sets_indices[i]):
                    print("\tCycle %d: arrays = %s" % (k + 1, str(v)))

        N_cycles = np.max(
            [len(slice_sets_indices[i]) for i in range(N_stacks)])

        reference = st.Stack.from_stack(self._reference)
        alphas = np.linspace(self._alpha_range[0], self._alpha_range[1],
                             N_cycles + 1)
        alphas = alphas[0:N_cycles]

        ctr_iter = [0]
        for i_cycle in range(0, N_cycles):
            self._registration_method.set_moving(reference)

            slice_index_sets_of_stacks = {
                i: (slice_sets_indices[i][i_cycle]
                    if i_cycle in slice_sets_indices[i] else [])
                for i in range(len(self._stacks))
            }

            ss2vreg = SliceSetToVolumeRegistration(
                print_prefix="Cycle %d/%d -- " % (i_cycle + 1, N_cycles),
                stacks=self._stacks,
                reference=reference,
                registration_method=self._registration_method,
                slice_index_sets_of_stacks=slice_index_sets_of_stacks,
                verbose=self._verbose,
            )
            ss2vreg.run()
            self._computational_time_registration += \
                ss2vreg.get_computational_time()

            # SRR step
            self._reconstruction_method.set_alpha(alphas[i_cycle])
            self._reconstruction_method.run()

            self._computational_time_reconstruction += \
                self._reconstruction_method.get_computational_time()

            reference = self._reconstruction_method.get_reconstruction()

            # Store SRR
            filename = "Iter%d_%s" % (
                ph.add_one(ctr_iter),
                self._reconstruction_method.get_setting_specific_filename())
            self._reconstructions.insert(
                0, st.Stack.from_stack(reference, filename=filename))
            if self._verbose:
                sitkh.show_stacks(self._reconstructions)

        # Run slice-to-volume registration in case last hierarchical run was
        # not based on individual slices
        if N_min > 1:
            s2vreg = SliceToVolumeRegistration(
                stacks=self._stacks,
                reference=reference,
                registration_method=self._registration_method,
                verbose=self._verbose)
            s2vreg.run()
            self._computational_time_registration += \
                s2vreg.get_computational_time()

            # SRR step
            self._reconstruction_method.set_alpha(alphas[-1])
            self._reconstruction_method.run()

            self._computational_time_reconstruction += \
                self._reconstruction_method.get_computational_time()

            # Store SRR
            filename = "Iter%d_%s" % (
                ph.add_one(ctr_iter),
                self._reconstruction_method.get_setting_specific_filename())
            self._reconstructions.insert(
                0, st.Stack.from_stack(reference, filename=filename))

            if self._verbose:
                sitkh.show_stacks(self._reconstructions)
    def _run(self):

        ph.print_title("Slice-to-Volume Registration")

        self._registration_method.set_moving(self._reference)

        for i, stack in enumerate(self._stacks):
            slices = stack.get_slices()
            transforms_sitk = [None] * len(slices)

            for j, slice_j in enumerate(slices):

                txt = "%sSlice-to-Volume Registration -- " \
                    "Stack %d/%d -- Slice %d/%d" % (
                        self._print_prefix,
                        i + 1, len(self._stacks),
                        j + 1, len(slices))
                if self._verbose:
                    ph.print_subtitle(txt)
                else:
                    ph.print_info(txt)

                self._registration_method.set_fixed(slice_j)
                self._registration_method.run()

                # Store information on registration transform
                transform_sitk = \
                    self._registration_method.\
                    get_registration_transform_sitk()
                transforms_sitk[j] = transform_sitk

            # Avoid slice misregistrations
            if self._s2v_smoothing is not None:
                ph.print_subtitle("Robust slice motion estimation "
                                  "(GP smoothing = %g, interleave = %d)" %
                                  (self._s2v_smoothing, self._interleave))
                robust_motion_estimator = rme.RobustMotionEstimator(
                    transforms_sitk=transforms_sitk,
                    interleave=self._interleave)
                robust_motion_estimator.run_gaussian_process_smoothing(
                    self._s2v_smoothing)
                transforms_sitk = \
                    robust_motion_estimator.get_robust_transforms_sitk()

                # Export figures
                # title = "%s_Stack%d%s" % (
                #     self._print_prefix, i, stack.get_filename())
                # title = ph.replace_string_for_print(title)
                # robust_motion_estimator.show_estimated_transform_parameters(
                #     dir_output="/tmp/fetal_brain/figs", title=title)

            # dir_output = "/tmp/fetal/figs"
            # motion_evaluator = me.MotionEvaluator(transforms_sitk)
            # motion_evaluator.run()
            # motion_evaluator.display(dir_output=dir_output, title=title)
            # motion_evaluator.show(dir_output=dir_output, title=title)

            # Update position of slice
            for j, slice_j in enumerate(slices):
                slice_j.update_motion_correction(transforms_sitk[j])

        # Reject misregistered slices
        if self._threshold is not None:
            ph.print_subtitle("Slice Outlier Rejection (%s < %g)" %
                              (self._threshold_measure, self._threshold))
            outlier_rejector = outre.OutlierRejector(
                stacks=self._stacks,
                reference=self._reference,
                threshold=self._threshold,
                measure=self._threshold_measure,
                verbose=True,
            )
            outlier_rejector.run()
            self._stacks = outlier_rejector.get_stacks()

            if len(self._stacks) == 0:
                raise RuntimeError(
                    "All slices of all stacks were rejected "
                    "as outliers. Volumetric reconstruction is aborted.")
Beispiel #10
0
def main():

    input_parser = InputArgparser(
        description="Script to export a side-by-side comparison of originally "
        "acquired and simulated/projected slice given the estimated "
        "volumetric reconstruction."
        "This function takes the result of "
        "simulate_stacks_from_reconstruction.py as input.", )
    input_parser.add_filenames(required=True)
    input_parser.add_dir_output(required=True)
    input_parser.add_option(
        option_string="--prefix-simulated",
        type=str,
        help="Specify the prefix of the simulated stacks to distinguish them "
        "from the original data.",
        default="Simulated_",
    )
    input_parser.add_option(
        option_string="--dir-input-simulated",
        type=str,
        help="Specify the directory where the simulated stacks are. "
        "If not given, it is assumed that they are in the same directory "
        "as the original ones.",
        default=None)
    input_parser.add_option(
        option_string="--resize",
        type=float,
        help="Factor to resize images (otherwise they might be very small "
        "depending on the FOV)",
        default=3)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    # Read original data
    filenames_original = args.filenames
    data_reader = dr.MultipleImagesReader(filenames_original)
    data_reader.read_data()
    stacks_original = data_reader.get_data()

    # Read data simulated from obtained reconstruction
    if args.dir_input_simulated is None:
        dir_input_simulated = os.path.dirname(filenames_original[0])
    else:
        dir_input_simulated = args.dir_input_simulated
    filenames_simulated = [
        os.path.join("%s", "%s%s") %
        (dir_input_simulated, args.prefix_simulated, os.path.basename(f))
        for f in filenames_original
    ]
    data_reader = dr.MultipleImagesReader(filenames_simulated)
    data_reader.read_data()
    stacks_simulated = data_reader.get_data()

    ph.create_directory(args.dir_output)

    for i in range(len(stacks_original)):
        try:
            stacks_original[i].sitk - stacks_simulated[i].sitk
        except:
            raise IOError(
                "Images '%s' and '%s' do not occupy the same space!" %
                (filenames_original[i], filenames_simulated[i]))

    # ---------------------Create side-by-side comparisons---------------------
    ph.print_title("Create side-by-side comparisons")
    intensity_max = 255
    intensity_min = 0
    for i in range(len(stacks_original)):
        ph.print_subtitle("Stack %d/%d" % (i + 1, len(stacks_original)))
        nda_3D_original = sitk.GetArrayFromImage(stacks_original[i].sitk)
        nda_3D_simulated = sitk.GetArrayFromImage(stacks_simulated[i].sitk)

        # Scale uniformly between 0 and 255 according to the simulated stack
        # for export to png
        scale = np.max(nda_3D_simulated)
        nda_3D_original = intensity_max * nda_3D_original / scale
        nda_3D_simulated = intensity_max * nda_3D_simulated / scale

        nda_3D_simulated = np.clip(nda_3D_simulated, intensity_min,
                                   intensity_max)
        nda_3D_original = np.clip(nda_3D_original, intensity_min,
                                  intensity_max)

        filename = stacks_original[i].get_filename()
        path_to_file = os.path.join(args.dir_output, "%s.pdf" % filename)

        # Export side-by-side comparison of each stack to a pdf file
        export_comparison_to_file(nda_3D_original,
                                  nda_3D_simulated,
                                  path_to_file,
                                  resize=args.resize)
Beispiel #11
0
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    # Read input
    input_parser = InputArgparser(
        description="Volumetric MRI reconstruction framework to reconstruct "
        "an isotropic, high-resolution 3D volume from multiple "
        "motion-corrected (or static) stacks of low-resolution slices.", )
    input_parser.add_dir_input()
    input_parser.add_filenames()
    input_parser.add_image_selection()
    input_parser.add_dir_output(required=True)
    input_parser.add_prefix_output(default="SRR_")
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_target_stack_index(default=0)
    input_parser.add_extra_frame_target(default=10)
    input_parser.add_isotropic_resolution(default=None)
    input_parser.add_reconstruction_space(default=None)
    input_parser.add_minimizer(default="lsmr")
    input_parser.add_iter_max(default=10)
    input_parser.add_reconstruction_type(default="TK1L2")
    input_parser.add_data_loss(default="linear")
    input_parser.add_data_loss_scale(default=1)
    input_parser.add_alpha(default=0.02  # TK1L2
                           # default=0.006  #TVL2, HuberL2
                           )
    input_parser.add_rho(default=0.5)
    input_parser.add_tv_solver(default="PD")
    input_parser.add_pd_alg_type(default="ALG2")
    input_parser.add_iterations(default=15)
    input_parser.add_subfolder_comparison()
    input_parser.add_provide_comparison(default=0)
    input_parser.add_log_script_execution(default=1)
    input_parser.add_verbose(default=0)
    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    # Write script execution call
    if args.log_script_execution:
        input_parser.write_performed_script_execution(
            os.path.abspath(__file__))

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    # Neither '--dir-input' nor '--filenames' was specified
    if args.filenames is not None and args.dir_input is not None:
        raise IOError("Provide input by either '--dir-input' or '--filenames' "
                      "but not both together")

    # '--dir-input' specified
    elif args.dir_input is not None:
        data_reader = dr.ImageSlicesDirectoryReader(
            path_to_directory=args.dir_input,
            suffix_mask=args.suffix_mask,
            image_selection=args.image_selection)

    # '--filenames' specified
    elif args.filenames is not None:
        data_reader = dr.MultipleImagesReader(args.filenames,
                                              suffix_mask=args.suffix_mask)

    else:
        raise IOError("Provide input by either '--dir-input' or '--filenames'")

    if args.reconstruction_type not in ["TK1L2", "TVL2", "HuberL2"]:
        raise IOError("Reconstruction type unknown")

    data_reader.read_data()
    stacks = data_reader.get_data()
    ph.print_info("%d input stacks read for further processing" % len(stacks))

    # Reconstruction space is given isotropically resampled target stack
    if args.reconstruction_space is None:
        recon0 = \
            stacks[args.target_stack_index].get_isotropically_resampled_stack(
                resolution=args.isotropic_resolution,
                extra_frame=args.extra_frame_target)

    # Reconstruction space was provided by user
    else:
        recon0 = st.Stack.from_filename(args.reconstruction_space,
                                        extract_slices=False)

        # Change resolution for isotropic resolution if provided by user
        if args.isotropic_resolution is not None:
            recon0 = recon0.get_isotropically_resampled_stack(
                args.isotropic_resolution)

        # Use image information of selected target stack as recon0 serves
        # as initial value for reconstruction
        recon0 = \
            stacks[args.target_stack_index].get_resampled_stack(recon0.sitk)
        recon0 = recon0.get_stack_multiplied_with_mask()

    if args.reconstruction_type in ["TVL2", "HuberL2"]:
        ph.print_title("Compute Initial value for %s" %
                       args.reconstruction_type)
    SRR0 = tk.TikhonovSolver(
        stacks=stacks,
        reconstruction=recon0,
        alpha=args.alpha,
        iter_max=args.iter_max,
        reg_type="TK1",
        minimizer=args.minimizer,
        data_loss=args.data_loss,
        data_loss_scale=args.data_loss_scale,
        # verbose=args.verbose,
    )
    SRR0.run()

    recon = SRR0.get_reconstruction()
    recon.set_filename(SRR0.get_setting_specific_filename(args.prefix_output))
    recon.write(args.dir_output)

    # List to store SRRs
    recons = []
    for i in range(0, len(stacks)):
        recons.append(stacks[i])
    recons.insert(0, recon)

    if args.reconstruction_type in ["TVL2", "HuberL2"]:
        ph.print_title("Compute %s reconstruction" % args.reconstruction_type)
        if args.tv_solver == "ADMM":
            SRR = admm.ADMMSolver(
                stacks=stacks,
                reconstruction=st.Stack.from_stack(SRR0.get_reconstruction()),
                minimizer=args.minimizer,
                alpha=args.alpha,
                iter_max=args.iter_max,
                rho=args.rho,
                data_loss=args.data_loss,
                iterations=args.iterations,
                verbose=args.verbose,
            )
            SRR.run()
            recon = SRR.get_reconstruction()
            recon.set_filename(
                SRR.get_setting_specific_filename(args.prefix_output))
            recons.insert(0, recon)

            recon.write(args.dir_output)

        else:

            SRR = pd.PrimalDualSolver(
                stacks=stacks,
                reconstruction=st.Stack.from_stack(SRR0.get_reconstruction()),
                minimizer=args.minimizer,
                alpha=args.alpha,
                iter_max=args.iter_max,
                iterations=args.iterations,
                alg_type=args.pd_alg_type,
                reg_type="TV"
                if args.reconstruction_type == "TVL2" else "huber",
                data_loss=args.data_loss,
                verbose=args.verbose,
            )
            SRR.run()
            recon = SRR.get_reconstruction()
            recon.set_filename(
                SRR.get_setting_specific_filename(args.prefix_output))
            recons.insert(0, recon)

            recon.write(args.dir_output)

    if args.verbose and not args.provide_comparison:
        sitkh.show_stacks(recons)

    # Show SRR together with linearly resampled input data.
    # Additionally, a script is generated to open files
    if args.provide_comparison:
        sitkh.show_stacks(
            recons,
            show_comparison_file=args.provide_comparison,
            dir_output=os.path.join(args.dir_output,
                                    args.subfolder_comparison),
        )

    ph.print_line_separator()

    elapsed_time = ph.stop_timing(time_start)
    ph.print_title("Summary")
    print("Computational Time for Volumetric Reconstruction: %s" %
          (elapsed_time))

    return 0
Beispiel #12
0
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    # Read input
    input_parser = InputArgparser(
        description="Volumetric MRI reconstruction framework to reconstruct "
        "an isotropic, high-resolution 3D volume from multiple stacks of 2D "
        "slices with motion correction. The resolution of the computed "
        "Super-Resolution Reconstruction (SRR) is given by the in-plane "
        "spacing of the selected target stack. A region of interest can be "
        "specified by providing a mask for the selected target stack. Only "
        "this region will then be reconstructed by the SRR algorithm which "
        "can substantially reduce the computational time.", )
    input_parser.add_dir_input()
    input_parser.add_filenames()
    input_parser.add_dir_output(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_target_stack_index(default=0)
    input_parser.add_search_angle(default=90)
    input_parser.add_multiresolution(default=0)
    input_parser.add_shrink_factors(default=[2, 1])
    input_parser.add_smoothing_sigmas(default=[1, 0])
    input_parser.add_sigma(default=0.9)
    input_parser.add_reconstruction_type(default="TK1L2")
    input_parser.add_iterations(default=15)
    input_parser.add_alpha(default=0.02)
    input_parser.add_alpha_first(default=0.05)
    input_parser.add_iter_max(default=10)
    input_parser.add_iter_max_first(default=5)
    input_parser.add_dilation_radius(default=3)
    input_parser.add_extra_frame_target(default=10)
    input_parser.add_bias_field_correction(default=0)
    input_parser.add_intensity_correction(default=0)
    input_parser.add_isotropic_resolution(default=None)
    input_parser.add_log_script_execution(default=1)
    input_parser.add_subfolder_motion_correction()
    input_parser.add_provide_comparison(default=0)
    input_parser.add_subfolder_comparison()
    input_parser.add_write_motion_correction(default=1)
    input_parser.add_verbose(default=0)
    input_parser.add_two_step_cycles(default=3)
    input_parser.add_use_masks_srr(default=1)
    input_parser.add_boundary_stacks(default=[10, 10, 0])
    input_parser.add_reference()
    input_parser.add_reference_mask()

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    # Write script execution call
    if args.log_script_execution:
        input_parser.write_performed_script_execution(
            os.path.abspath(__file__))

    # Use FLIRT for volume-to-volume reg. step. Otherwise, RegAladin is used.
    use_flirt_for_v2v_registration = True

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    # Neither '--dir-input' nor '--filenames' was specified
    if args.filenames is not None and args.dir_input is not None:
        raise IOError("Provide input by either '--dir-input' or '--filenames' "
                      "but not both together")

    # '--dir-input' specified
    elif args.dir_input is not None:
        data_reader = dr.ImageDirectoryReader(args.dir_input,
                                              suffix_mask=args.suffix_mask)

    # '--filenames' specified
    elif args.filenames is not None:
        data_reader = dr.MultipleImagesReader(args.filenames,
                                              suffix_mask=args.suffix_mask)

    else:
        raise IOError("Provide input by either '--dir-input' or '--filenames'")

    if len(args.boundary_stacks) is not 3:
        raise IOError(
            "Provide exactly three values for '--boundary-stacks' to define "
            "cropping in i-, j-, and k-dimension of the input stacks")

    data_reader.read_data()
    stacks = data_reader.get_data()
    ph.print_info("%d input stacks read for further processing" % len(stacks))

    if all(s.is_unity_mask() is True for s in stacks):
        ph.print_warning("No mask is provided! "
                         "Generated reconstruction space may be very big!")

    # ---------------------------Data Preprocessing---------------------------
    ph.print_title("Data Preprocessing")

    segmentation_propagator = segprop.SegmentationPropagation(
        # registration_method=regflirt.FLIRT(use_verbose=args.verbose),
        dilation_radius=args.dilation_radius,
        dilation_kernel="Ball",
    )

    data_preprocessing = dp.DataPreprocessing(
        stacks=stacks,
        segmentation_propagator=segmentation_propagator,
        use_cropping_to_mask=True,
        use_N4BiasFieldCorrector=args.bias_field_correction,
        use_intensity_correction=args.intensity_correction,
        target_stack_index=args.target_stack_index,
        boundary_i=args.boundary_stacks[0],
        boundary_j=args.boundary_stacks[1],
        boundary_k=args.boundary_stacks[2],
        unit="mm",
    )
    data_preprocessing.run()
    time_data_preprocessing = data_preprocessing.get_computational_time()

    # Get preprocessed stacks
    stacks = data_preprocessing.get_preprocessed_stacks()

    # Define reference/target stack for registration and reconstruction
    if args.reference is not None:
        reference = st.Stack.from_filename(file_path=args.reference,
                                           file_path_mask=args.reference_mask,
                                           extract_slices=False)

    else:
        reference = st.Stack.from_stack(stacks[args.target_stack_index])

    # ------------------------Volume-to-Volume Registration--------------------
    if args.two_step_cycles > 0:
        # Define search angle ranges for FLIRT in all three dimensions
        search_angles = [
            "-searchr%s -%d %d" % (x, args.search_angle, args.search_angle)
            for x in ["x", "y", "z"]
        ]
        search_angles = (" ").join(search_angles)

        if use_flirt_for_v2v_registration:
            vol_registration = regflirt.FLIRT(
                registration_type="Rigid",
                use_fixed_mask=True,
                use_moving_mask=True,
                options=search_angles,
                use_verbose=False,
            )
        else:
            vol_registration = niftyreg.RegAladin(
                registration_type="Rigid",
                use_fixed_mask=True,
                use_moving_mask=True,
                use_verbose=False,
            )
        v2vreg = pipeline.VolumeToVolumeRegistration(
            stacks=stacks,
            reference=reference,
            registration_method=vol_registration,
            verbose=args.verbose,
        )
        v2vreg.run()
        stacks = v2vreg.get_stacks()
        time_registration = v2vreg.get_computational_time()

    else:
        time_registration = ph.get_zero_time()

    # ---------------------------Create first volume---------------------------
    time_tmp = ph.start_timing()

    # Isotropic resampling to define HR target space
    ph.print_title("Reconstruction Space Generation")
    HR_volume = reference.get_isotropically_resampled_stack(
        resolution=args.isotropic_resolution)
    ph.print_info(
        "Isotropic reconstruction space with %g mm resolution is created" %
        HR_volume.sitk.GetSpacing()[0])

    if args.reference is None:
        # Create joint image mask in target space
        joint_image_mask_builder = imb.JointImageMaskBuilder(
            stacks=stacks,
            target=HR_volume,
            dilation_radius=1,
        )
        joint_image_mask_builder.run()
        HR_volume = joint_image_mask_builder.get_stack()
        ph.print_info("Isotropic reconstruction space is centered around "
                      "joint stack masks. ")

        # Crop to space defined by mask (plus extra margin)
        HR_volume = HR_volume.get_cropped_stack_based_on_mask(
            boundary_i=args.extra_frame_target,
            boundary_j=args.extra_frame_target,
            boundary_k=args.extra_frame_target,
            unit="mm",
        )

        # Scattered Data Approximation to get first estimate of HR volume
        ph.print_title("First Estimate of HR Volume")
        SDA = sda.ScatteredDataApproximation(stacks,
                                             HR_volume,
                                             sigma=args.sigma)
        SDA.run()
        HR_volume = SDA.get_reconstruction()

    time_reconstruction = ph.stop_timing(time_tmp)

    if args.verbose:
        tmp = list(stacks)
        tmp.insert(0, HR_volume)
        sitkh.show_stacks(tmp, segmentation=HR_volume)

    # ----------------Two-step Slice-to-Volume Registration SRR----------------
    SRR = tk.TikhonovSolver(
        stacks=stacks,
        reconstruction=HR_volume,
        reg_type="TK1",
        minimizer="lsmr",
        alpha=args.alpha_first,
        iter_max=args.iter_max_first,
        verbose=True,
        use_masks=args.use_masks_srr,
    )

    if args.two_step_cycles > 0:

        registration = regsitk.SimpleItkRegistration(
            moving=HR_volume,
            use_fixed_mask=True,
            use_moving_mask=True,
            use_verbose=args.verbose,
            interpolator="Linear",
            metric="Correlation",
            use_multiresolution_framework=args.multiresolution,
            shrink_factors=args.shrink_factors,
            smoothing_sigmas=args.smoothing_sigmas,
            initializer_type="SelfGEOMETRY",
            optimizer="ConjugateGradientLineSearch",
            optimizer_params={
                "learningRate": 1,
                "numberOfIterations": 100,
                "lineSearchUpperLimit": 2,
            },
            scales_estimator="Jacobian",
        )
        two_step_s2v_reg_recon = \
            pipeline.TwoStepSliceToVolumeRegistrationReconstruction(
                stacks=stacks,
                reference=HR_volume,
                registration_method=registration,
                reconstruction_method=SRR,
                cycles=args.two_step_cycles,
                alpha_range=[args.alpha_first, args.alpha],
                verbose=args.verbose,
            )
        two_step_s2v_reg_recon.run()
        HR_volume_iterations = \
            two_step_s2v_reg_recon.get_iterative_reconstructions()
        time_registration += \
            two_step_s2v_reg_recon.get_computational_time_registration()
        time_reconstruction += \
            two_step_s2v_reg_recon.get_computational_time_reconstruction()
    else:
        HR_volume_iterations = []

    # Write motion-correction results
    if args.write_motion_correction:
        for stack in stacks:
            stack.write(
                os.path.join(args.dir_output,
                             args.subfolder_motion_correction),
                write_mask=True,
                write_slices=True,
                write_transforms=True,
                suffix_mask=args.suffix_mask,
            )

    # ------------------Final Super-Resolution Reconstruction------------------
    ph.print_title("Final Super-Resolution Reconstruction")
    if args.reconstruction_type in ["TVL2", "HuberL2"]:
        SRR = pd.PrimalDualSolver(
            stacks=stacks,
            reconstruction=HR_volume,
            reg_type="TV" if args.reconstruction_type == "TVL2" else "huber",
            iterations=args.iterations,
        )
    else:
        SRR = tk.TikhonovSolver(
            stacks=stacks,
            reconstruction=HR_volume,
            reg_type="TK1" if args.reconstruction_type == "TK1L2" else "TK0",
            use_masks=args.use_masks_srr,
        )
    SRR.set_alpha(args.alpha)
    SRR.set_iter_max(args.iter_max)
    SRR.set_verbose(True)
    SRR.run()
    time_reconstruction += SRR.get_computational_time()

    elapsed_time_total = ph.stop_timing(time_start)

    # Write SRR result
    HR_volume_final = SRR.get_reconstruction()
    HR_volume_final.set_filename(SRR.get_setting_specific_filename())
    HR_volume_final.write(args.dir_output,
                          write_mask=True,
                          suffix_mask=args.suffix_mask)

    HR_volume_iterations.insert(0, HR_volume_final)
    for stack in stacks:
        HR_volume_iterations.append(stack)

    if args.verbose and not args.provide_comparison:
        sitkh.show_stacks(HR_volume_iterations, segmentation=HR_volume)
    # HR_volume_final.show()

    # Show SRR together with linearly resampled input data.
    # Additionally, a script is generated to open files
    if args.provide_comparison:
        sitkh.show_stacks(
            HR_volume_iterations,
            segmentation=HR_volume,
            show_comparison_file=args.provide_comparison,
            dir_output=os.path.join(args.dir_output,
                                    args.subfolder_comparison),
        )

    # Summary
    ph.print_title("Summary")
    print("Computational Time for Data Preprocessing: %s" %
          (time_data_preprocessing))
    print("Computational Time for Registrations: %s" % (time_registration))
    print("Computational Time for Reconstructions: %s" % (time_reconstruction))
    print("Computational Time for Entire Reconstruction Pipeline: %s" %
          (elapsed_time_total))

    ph.print_line_separator()

    return 0
Beispiel #13
0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Register an obtained reconstruction (moving) "
        "to a template image/space (fixed) using rigid registration. "
        "The resulting registration can optionally be applied to previously "
        "obtained motion correction slice transforms so that a volumetric "
        "reconstruction is possible in the (standard anatomical) space "
        "defined by the fixed.", )
    input_parser.add_fixed(required=True)
    input_parser.add_moving(required=True)
    input_parser.add_output(help="Path to registration transform (.txt)",
                            required=True)
    input_parser.add_fixed_mask()
    input_parser.add_moving_mask()
    input_parser.add_dir_input_mc()
    input_parser.add_search_angle(default=180)
    input_parser.add_option(option_string="--initial-transform",
                            type=str,
                            help="Path to initial transform.",
                            default=None)
    input_parser.add_option(
        option_string="--test-ap-flip",
        type=int,
        help="Turn on/off functionality to run an additional registration "
        "after an AP-flip. Seems to be more robust to find a better "
        "registration outcome in general.",
        default=1)
    input_parser.add_option(
        option_string="--use-flirt",
        type=int,
        help="Turn on/off functionality to use FLIRT for the registration.",
        default=1)
    input_parser.add_option(
        option_string="--use-regaladin",
        type=int,
        help="Turn on/off functionality to use RegAladin for the "
        "registration.",
        default=1)
    input_parser.add_verbose(default=0)
    input_parser.add_log_config(default=1)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    debug = 0

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    if not args.use_regaladin and not args.use_flirt:
        raise IOError("Either RegAladin or FLIRT must be activated.")

    if not args.output.endswith(".txt"):
        raise IOError("output transformation path must end in '.txt'")

    dir_output = os.path.dirname(args.output)

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")
    fixed = st.Stack.from_filename(file_path=args.fixed,
                                   file_path_mask=args.fixed_mask,
                                   extract_slices=False)
    moving = st.Stack.from_filename(file_path=args.moving,
                                    file_path_mask=args.moving_mask,
                                    extract_slices=False)

    if args.initial_transform is not None:
        transform_sitk = sitkh.read_transform_sitk(args.initial_transform)
    else:
        transform_sitk = sitk.AffineTransform(fixed.sitk.GetDimension())
    sitk.WriteTransform(transform_sitk, args.output)

    path_to_tmp_output = os.path.join(
        DIR_TMP, ph.append_to_filename(os.path.basename(args.moving),
                                       "_warped"))

    # -------------------Register Reconstruction to Template-------------------
    ph.print_title("Register Reconstruction to Template")

    if args.use_flirt:
        path_to_transform_flirt = os.path.join(DIR_TMP, "transform_flirt.txt")

        # Convert SimpleITK into FLIRT transform
        cmd = "simplereg_transform -sitk2flirt %s %s %s %s" % (
            args.output, args.fixed, args.moving, path_to_transform_flirt)
        ph.execute_command(cmd, verbose=False)

        # Define search angle ranges for FLIRT in all three dimensions
        search_angles = [
            "-searchr%s -%d %d" % (x, args.search_angle, args.search_angle)
            for x in ["x", "y", "z"]
        ]

        # flt = nipype.interfaces.fsl.FLIRT()
        # flt.inputs.in_file = args.moving
        # flt.inputs.reference = args.fixed
        # if args.initial_transform is not None:
        #     flt.inputs.in_matrix_file = path_to_transform_flirt
        # flt.inputs.out_matrix_file = path_to_transform_flirt
        # # flt.inputs.output_type = "NIFTI_GZ"
        # flt.inputs.out_file = path_to_tmp_output
        # flt.inputs.args = "-dof 6"
        # flt.inputs.args += " %s" % " ".join(search_angles)
        # if args.moving_mask is not None:
        #     flt.inputs.in_weight = args.moving_mask
        # if args.fixed_mask is not None:
        #     flt.inputs.ref_weight = args.fixed_mask
        # ph.print_info("Run Registration (FLIRT) ... ", newline=False)
        # flt.run()
        # print("done")

        cmd_args = ["flirt"]
        cmd_args.append("-in %s" % args.moving)
        cmd_args.append("-ref %s" % args.fixed)
        if args.initial_transform is not None:
            cmd_args.append("-init %s" % path_to_transform_flirt)
        cmd_args.append("-omat %s" % path_to_transform_flirt)
        cmd_args.append("-out %s" % path_to_tmp_output)
        cmd_args.append("-dof 6")
        cmd_args.append((" ").join(search_angles))
        if args.moving_mask is not None:
            cmd_args.append("-inweight %s" % args.moving_mask)
        if args.fixed_mask is not None:
            cmd_args.append("-refweight %s" % args.fixed_mask)
        ph.print_info("Run Registration (FLIRT) ... ", newline=False)
        ph.execute_command(" ".join(cmd_args), verbose=False)
        print("done")

        # Convert FLIRT to SimpleITK transform
        cmd = "simplereg_transform -flirt2sitk %s %s %s %s" % (
            path_to_transform_flirt, args.fixed, args.moving, args.output)
        ph.execute_command(cmd, verbose=False)

        if debug:
            ph.show_niftis([args.fixed, path_to_tmp_output])

    # Additionally, use RegAladin for more accurate alignment
    # Rationale: FLIRT has better capture range, but RegAladin seems to
    # find better alignment once it is within its capture range.
    if args.use_regaladin:
        path_to_transform_regaladin = os.path.join(DIR_TMP,
                                                   "transform_regaladin.txt")

        # Convert SimpleITK to RegAladin transform
        cmd = "simplereg_transform -sitk2nreg %s %s" % (
            args.output, path_to_transform_regaladin)
        ph.execute_command(cmd, verbose=False)

        # nreg = nipype.interfaces.niftyreg.RegAladin()
        # nreg.inputs.ref_file = args.fixed
        # nreg.inputs.flo_file = args.moving
        # nreg.inputs.res_file = path_to_tmp_output
        # nreg.inputs.in_aff_file = path_to_transform_regaladin
        # nreg.inputs.aff_file = path_to_transform_regaladin
        # nreg.inputs.args = "-rigOnly -voff"
        # if args.moving_mask is not None:
        #     nreg.inputs.fmask_file = args.moving_mask
        # if args.fixed_mask is not None:
        #     nreg.inputs.rmask_file = args.fixed_mask
        # ph.print_info("Run Registration (RegAladin) ... ", newline=False)
        # nreg.run()
        # print("done")

        cmd_args = ["reg_aladin"]
        cmd_args.append("-ref %s" % args.fixed)
        cmd_args.append("-flo %s" % args.moving)
        cmd_args.append("-res %s" % path_to_tmp_output)
        if args.initial_transform is not None or args.use_flirt == 1:
            cmd_args.append("-inaff %s" % path_to_transform_regaladin)
        cmd_args.append("-aff %s" % path_to_transform_regaladin)
        # cmd_args.append("-cog")
        # cmd_args.append("-ln 2")
        cmd_args.append("-rigOnly")
        cmd_args.append("-voff")
        if args.moving_mask is not None:
            cmd_args.append("-fmask %s" % args.moving_mask)
        if args.fixed_mask is not None:
            cmd_args.append("-rmask %s" % args.fixed_mask)
        ph.print_info("Run Registration (RegAladin) ... ", newline=False)
        ph.execute_command(" ".join(cmd_args), verbose=False)
        print("done")

        # Convert RegAladin to SimpleITK transform
        cmd = "simplereg_transform -nreg2sitk %s %s" % (
            path_to_transform_regaladin, args.output)
        ph.execute_command(cmd, verbose=False)

        if debug:
            ph.show_niftis([args.fixed, path_to_tmp_output])

    if args.test_ap_flip:
        path_to_transform_flip = os.path.join(DIR_TMP, "transform_flip.txt")
        path_to_tmp_output_flip = os.path.join(DIR_TMP, "output_flip.nii.gz")

        # Get AP-flip transform
        transform_ap_flip_sitk = get_ap_flip_transform(args.fixed)
        path_to_transform_flip_regaladin = os.path.join(
            DIR_TMP, "transform_flip_regaladin.txt")
        sitk.WriteTransform(transform_ap_flip_sitk, path_to_transform_flip)

        # Compose current transform with AP flip transform
        cmd = "simplereg_transform -c %s %s %s" % (
            args.output, path_to_transform_flip, path_to_transform_flip)
        ph.execute_command(cmd, verbose=False)

        # Convert SimpleITK to RegAladin transform
        cmd = "simplereg_transform -sitk2nreg %s %s" % (
            path_to_transform_flip, path_to_transform_flip_regaladin)
        ph.execute_command(cmd, verbose=False)

        # nreg = nipype.interfaces.niftyreg.RegAladin()
        # nreg.inputs.ref_file = args.fixed
        # nreg.inputs.flo_file = args.moving
        # nreg.inputs.res_file = path_to_tmp_output_flip
        # nreg.inputs.in_aff_file = path_to_transform_flip_regaladin
        # nreg.inputs.aff_file = path_to_transform_flip_regaladin
        # nreg.inputs.args = "-rigOnly -voff"
        # if args.moving_mask is not None:
        #     nreg.inputs.fmask_file = args.moving_mask
        # if args.fixed_mask is not None:
        #     nreg.inputs.rmask_file = args.fixed_mask
        # ph.print_info("Run Registration AP-flipped (RegAladin) ... ",
        #               newline=False)
        # nreg.run()
        # print("done")

        cmd_args = ["reg_aladin"]
        cmd_args.append("-ref %s" % args.fixed)
        cmd_args.append("-flo %s" % args.moving)
        cmd_args.append("-res %s" % path_to_tmp_output_flip)
        cmd_args.append("-inaff %s" % path_to_transform_flip_regaladin)
        cmd_args.append("-aff %s" % path_to_transform_flip_regaladin)
        cmd_args.append("-rigOnly")
        # cmd_args.append("-ln 2")
        cmd_args.append("-voff")
        if args.moving_mask is not None:
            cmd_args.append("-fmask %s" % args.moving_mask)
        if args.fixed_mask is not None:
            cmd_args.append("-rmask %s" % args.fixed_mask)
        ph.print_info("Run Registration AP-flipped (RegAladin) ... ",
                      newline=False)
        ph.execute_command(" ".join(cmd_args), verbose=False)
        print("done")

        if debug:
            ph.show_niftis(
                [args.fixed, path_to_tmp_output, path_to_tmp_output_flip])

        warped_moving = st.Stack.from_filename(path_to_tmp_output,
                                               extract_slices=False)
        warped_moving_flip = st.Stack.from_filename(path_to_tmp_output_flip,
                                                    extract_slices=False)
        fixed = st.Stack.from_filename(args.fixed, args.fixed_mask)

        stacks = [warped_moving, warped_moving_flip]
        image_similarity_evaluator = ise.ImageSimilarityEvaluator(
            stacks=stacks, reference=fixed)
        image_similarity_evaluator.compute_similarities()
        similarities = image_similarity_evaluator.get_similarities()

        if similarities["NMI"][1] > similarities["NMI"][0]:
            ph.print_info("AP-flipped outcome better")

            # Convert RegAladin to SimpleITK transform
            cmd = "simplereg_transform -nreg2sitk %s %s" % (
                path_to_transform_flip_regaladin, args.output)
            ph.execute_command(cmd, verbose=False)

            # Copy better outcome
            cmd = "cp -p %s %s" % (path_to_tmp_output_flip, path_to_tmp_output)
            ph.execute_command(cmd, verbose=False)

        else:
            ph.print_info("AP-flip does not improve outcome")

    if args.dir_input_mc is not None:
        transform_sitk = sitkh.read_transform_sitk(args.output, inverse=1)

        if args.dir_input_mc.endswith("/"):
            subdir_mc = args.dir_input_mc.split("/")[-2]
        else:
            subdir_mc = args.dir_input_mc.split("/")[-1]
        dir_output_mc = os.path.join(dir_output, subdir_mc)

        ph.create_directory(dir_output_mc, delete_files=True)
        pattern = REGEX_FILENAMES + "[.]tfm"
        p = re.compile(pattern)
        trafos = [t for t in os.listdir(args.dir_input_mc) if p.match(t)]
        for t in trafos:
            path_to_input_transform = os.path.join(args.dir_input_mc, t)
            path_to_output_transform = os.path.join(dir_output_mc, t)
            t_sitk = sitkh.read_transform_sitk(path_to_input_transform)
            t_sitk = sitkh.get_composite_sitk_affine_transform(
                transform_sitk, t_sitk)
            sitk.WriteTransform(t_sitk, path_to_output_transform)

    if args.verbose:
        ph.show_niftis([args.fixed, path_to_tmp_output])

    elapsed_time_total = ph.stop_timing(time_start)

    # Summary
    ph.print_title("Summary")
    print("Computational Time: %s" % (elapsed_time_total))

    return 0
Beispiel #14
0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Perform Bias Field correction using N4ITK.", )
    input_parser.add_filename(required=True)
    input_parser.add_output(required=True)
    input_parser.add_filename_mask()
    input_parser.add_option(
        option_string="--convergence-threshold",
        type=float,
        help="Specify the convergence threshold.",
        default=1e-6,
    )
    input_parser.add_option(
        option_string="--spline-order",
        type=int,
        help="Specify the spline order defining the bias field estimate.",
        default=3,
    )
    input_parser.add_option(
        option_string="--wiener-filter-noise",
        type=float,
        help="Specify the noise estimate defining the Wiener filter.",
        default=0.11,
    )
    input_parser.add_option(
        option_string="--bias-field-fwhm",
        type=float,
        help="Specify the full width at half maximum parameter characterizing "
        "the width of the Gaussian deconvolution.",
        default=0.15,
    )
    input_parser.add_log_config(default=1)
    input_parser.add_verbose(default=0)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if np.alltrue([not args.output.endswith(t) for t in ALLOWED_EXTENSIONS]):
        raise ValueError(
            "output filename invalid; allowed extensions are: %s" %
            ", ".join(ALLOWED_EXTENSIONS))

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    # Read data
    stack = st.Stack.from_filename(
        file_path=args.filename,
        file_path_mask=args.filename_mask,
        extract_slices=False,
    )

    # Perform Bias Field Correction
    # ph.print_title("Perform Bias Field Correction")
    bias_field_corrector = n4itk.N4BiasFieldCorrection(
        stack=stack,
        use_mask=True if args.filename_mask is not None else False,
        convergence_threshold=args.convergence_threshold,
        spline_order=args.spline_order,
        wiener_filter_noise=args.wiener_filter_noise,
        bias_field_fwhm=args.bias_field_fwhm,
    )
    ph.print_info("N4ITK Bias Field Correction ... ", newline=False)
    bias_field_corrector.run_bias_field_correction()
    stack_corrected = bias_field_corrector.get_bias_field_corrected_stack()
    print("done")

    dw.DataWriter.write_image(stack_corrected.sitk, args.output)

    elapsed_time = ph.stop_timing(time_start)

    if args.verbose:
        ph.show_niftis([args.filename, args.output])

    ph.print_title("Summary")
    exe_file_info = os.path.basename(os.path.abspath(__file__)).split(".")[0]
    print("%s | Computational Time for Bias Field Correction: %s" %
          (exe_file_info, elapsed_time))

    return 0
Beispiel #15
0
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    # Read input
    input_parser = InputArgparser(
        description="Script to study reconstruction parameters and their "
        "impact on the volumetric reconstruction quality.",
    )
    input_parser.add_dir_input()
    input_parser.add_filenames()
    input_parser.add_image_selection()
    input_parser.add_dir_output(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_reconstruction_space()
    input_parser.add_reference(
        help="Path to reference NIfTI image file. If given the volumetric "
        "reconstructed is performed in this physical space. "
        "Either a reconstruction space or a reference must be provided",
        required=False)
    input_parser.add_reference_mask(default=None)
    input_parser.add_study_name()
    input_parser.add_reconstruction_type(default="TK1L2")
    input_parser.add_measures(default=["PSNR", "RMSE", "SSIM", "NCC", "NMI"])
    input_parser.add_tv_solver(default="PD")
    input_parser.add_iterations(default=50)
    input_parser.add_rho(default=0.1)
    input_parser.add_iter_max(default=10)
    input_parser.add_minimizer(default="lsmr")
    input_parser.add_alpha(default=0.01)
    input_parser.add_data_loss(default="linear")
    input_parser.add_data_loss_scale(default=1)
    input_parser.add_log_script_execution(default=1)
    input_parser.add_verbose(default=1)

    # Range for parameter sweeps
    input_parser.add_alpha_range(default=[0.001, 0.05, 20])  # TK1L2
    # input_parser.add_alpha_range(default=[0.001, 0.003, 10])  # TVL2, HuberL2
    input_parser.add_data_losses(
        # default=["linear", "arctan"]
    )
    input_parser.add_data_loss_scale_range(
        # default=[0.1, 1.5, 2]
    )

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.reference is None and args.reconstruction_space is None:
        raise IOError("Either reference (--reference) or reconstruction space "
                      "(--reconstruction-space) must be provided.")

    # Write script execution call
    if args.log_script_execution:
        input_parser.write_performed_script_execution(
            os.path.abspath(__file__))

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    # Neither '--dir-input' nor '--filenames' was specified
    if args.filenames is not None and args.dir_input is not None:
        raise IOError(
            "Provide input by either '--dir-input' or '--filenames' "
            "but not both together")

    # '--dir-input' specified
    elif args.dir_input is not None:
        data_reader = dr.ImageSlicesDirectoryReader(
            path_to_directory=args.dir_input,
            suffix_mask=args.suffix_mask,
            image_selection=args.image_selection)

    # '--filenames' specified
    elif args.filenames is not None:
        data_reader = dr.MultipleImagesReader(
            args.filenames, suffix_mask=args.suffix_mask)

    else:
        raise IOError(
            "Provide input by either '--dir-input' or '--filenames'")

    data_reader.read_data()
    stacks = data_reader.get_data()
    ph.print_info("%d input stacks read for further processing" % len(stacks))

    if args.reference is not None:
        reference = st.Stack.from_filename(
            file_path=args.reference,
            file_path_mask=args.reference_mask,
            extract_slices=False)

        reconstruction_space = stacks[0].get_resampled_stack(reference.sitk)
        reconstruction_space = \
            reconstruction_space.get_stack_multiplied_with_mask()
        x_ref = sitk.GetArrayFromImage(reference.sitk).flatten()
        x_ref_mask = sitk.GetArrayFromImage(reference.sitk_mask).flatten()

    else:
        reconstruction_space = st.Stack.from_filename(
            file_path=args.reconstruction_space,
            extract_slices=False)
        reconstruction_space = stacks[0].get_resampled_stack(
            reconstruction_space.sitk)
        reconstruction_space = \
            reconstruction_space.get_stack_multiplied_with_mask()
        x_ref = None
        x_ref_mask = None

    # ----------------------------Set Up Parameters----------------------------
    parameters = {}
    parameters["alpha"] = np.linspace(
        args.alpha_range[0], args.alpha_range[1], int(args.alpha_range[2]))
    if args.data_losses is not None:
        parameters["data_loss"] = args.data_losses
    if args.data_loss_scale_range is not None:
        parameters["data_loss_scale"] = np.linspace(
            args.data_loss_scale_range[0],
            args.data_loss_scale_range[1],
            int(args.data_loss_scale_range[2]))

    # --------------------------Set Up Parameter Study-------------------------
    if args.study_name is None:
        name = args.reconstruction_type
    else:
        name = args.study_name

    reconstruction_info = {
        "shape": reconstruction_space.sitk.GetSize()[::-1],
        "origin": reconstruction_space.sitk.GetOrigin(),
        "spacing": reconstruction_space.sitk.GetSpacing(),
        "direction": reconstruction_space.sitk.GetDirection(),
    }

    # Create Tikhonov solver from which all information can be extracted
    # (also for other reconstruction types)
    tmp = tk.TikhonovSolver(
        stacks=stacks,
        reconstruction=reconstruction_space,
        alpha=args.alpha,
        iter_max=args.iter_max,
        data_loss=args.data_loss,
        data_loss_scale=args.data_loss_scale,
        reg_type="TK1",
        minimizer=args.minimizer,
        verbose=args.verbose,
    )
    solver = tmp.get_solver()

    parameter_study_interface = \
        deconv_interface.DeconvolutionParameterStudyInterface(
            A=solver.get_A(),
            A_adj=solver.get_A_adj(),
            D=solver.get_B(),
            D_adj=solver.get_B_adj(),
            b=solver.get_b(),
            x0=solver.get_x0(),
            alpha=solver.get_alpha(),
            x_scale=solver.get_x_scale(),
            data_loss=solver.get_data_loss(),
            data_loss_scale=solver.get_data_loss_scale(),
            iter_max=solver.get_iter_max(),
            minimizer=solver.get_minimizer(),
            iterations=args.iterations,
            measures=args.measures,
            dimension=3,
            L2=16./reconstruction_space.sitk.GetSpacing()[0]**2,
            reconstruction_type=args.reconstruction_type,
            rho=args.rho,
            dir_output=args.dir_output,
            parameters=parameters,
            name=name,
            reconstruction_info=reconstruction_info,
            x_ref=x_ref,
            x_ref_mask=x_ref_mask,
            tv_solver=args.tv_solver,
            verbose=args.verbose,
        )
    parameter_study_interface.set_up_parameter_study()
    parameter_study = parameter_study_interface.get_parameter_study()

    # Run parameter study
    parameter_study.run()

    print("\nComputational time for Deconvolution Parameter Study %s: %s" %
          (name, parameter_study.get_computational_time()))

    return 0
Beispiel #16
0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Perform Bias Field correction on images based on N4ITK.",
    )
    input_parser.add_filenames(required=True)
    input_parser.add_dir_output(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_prefix_output(default="N4ITK_")
    input_parser.add_option(
        option_string="--convergence-threshold",
        type=float,
        help="Specify the convergence threshold.",
        default=1e-6,
    )
    input_parser.add_option(
        option_string="--spline-order",
        type=int,
        help="Specify the spline order defining the bias field estimate.",
        default=3,
    )
    input_parser.add_option(
        option_string="--wiener-filter-noise",
        type=float,
        help="Specify the noise estimate defining the Wiener filter.",
        default=0.11,
    )
    input_parser.add_option(
        option_string="--bias-field-fwhm",
        type=float,
        help="Specify the full width at half maximum parameter characterizing "
        "the width of the Gaussian deconvolution.",
        default=0.15,
    )
    input_parser.add_log_script_execution(default=1)
    input_parser.add_verbose(default=0)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    # Write script execution call
    if args.log_script_execution:
        input_parser.write_performed_script_execution(
            os.path.abspath(__file__))

    # Read data
    data_reader = dr.MultipleImagesReader(args.filenames,
                                          suffix_mask=args.suffix_mask)
    data_reader.read_data()
    stacks = data_reader.get_data()

    # Perform Bias Field Correction
    ph.print_title("Perform Bias Field Correction")
    bias_field_corrector = n4itk.N4BiasFieldCorrection(
        convergence_threshold=args.convergence_threshold,
        spline_order=args.spline_order,
        wiener_filter_noise=args.wiener_filter_noise,
        bias_field_fwhm=args.bias_field_fwhm,
        prefix_corrected=args.prefix_output,
    )
    stacks_corrected = [None] * len(stacks)
    for i, stack in enumerate(stacks):
        ph.print_info("Image %d/%d: N4ITK Bias Field Correction ... " %
                      (i + 1, len(stacks)),
                      newline=False)
        bias_field_corrector.set_stack(stack)
        bias_field_corrector.run_bias_field_correction()
        stacks_corrected[i] = \
            bias_field_corrector.get_bias_field_corrected_stack()
        print("done")
        ph.print_info("Image %d/%d: Computational time = %s" %
                      (i + 1, len(stacks),
                       bias_field_corrector.get_computational_time()))

        # Write Data
        stacks_corrected[i].write(args.dir_output,
                                  write_mask=True,
                                  suffix_mask=args.suffix_mask)

        if args.verbose:
            sitkh.show_stacks([stacks[i], stacks_corrected[i]],
                              segmentation=stacks[i])

    elapsed_time = ph.stop_timing(time_start)

    ph.print_title("Summary")
    print("Computational Time for Bias Field Correction(s): %s" %
          (elapsed_time))

    return 0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Run reconstruction pipeline including "
        "(i) bias field correction, "
        "(ii) volumetric reconstruction in subject space, "
        "and (iii) volumetric reconstruction in template space.", )
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks(required=True)
    input_parser.add_target_stack(required=False)
    input_parser.add_suffix_mask(default="''")
    input_parser.add_dir_output(required=True)
    input_parser.add_alpha(default=0.01)
    input_parser.add_verbose(default=0)
    input_parser.add_gestational_age(required=False)
    input_parser.add_prefix_output(default="")
    input_parser.add_search_angle(default=180)
    input_parser.add_multiresolution(default=0)
    input_parser.add_log_config(default=1)
    input_parser.add_isotropic_resolution()
    input_parser.add_reference()
    input_parser.add_reference_mask()
    input_parser.add_bias_field_correction(default=1)
    input_parser.add_intensity_correction(default=1)
    input_parser.add_iter_max(default=10)
    input_parser.add_two_step_cycles(default=3)
    input_parser.add_option(
        option_string="--run-bias-field-correction",
        type=int,
        help="Turn on/off bias field correction. "
        "If off, it is assumed that this step was already performed",
        default=1)
    input_parser.add_option(
        option_string="--run-recon-subject-space",
        type=int,
        help="Turn on/off reconstruction in subject space. "
        "If off, it is assumed that this step was already performed",
        default=1)
    input_parser.add_option(
        option_string="--run-recon-template-space",
        type=int,
        help="Turn on/off reconstruction in template space. "
        "If off, it is assumed that this step was already performed",
        default=1)
    input_parser.add_option(
        option_string="--run-data-vs-simulated-data",
        type=int,
        help="Turn on/off comparison of data vs data simulated from the "
        "obtained volumetric reconstruction. "
        "If off, it is assumed that this step was already performed",
        default=0)
    input_parser.add_option(
        option_string="--initial-transform",
        type=str,
        help="Set initial transform to be used for register_image.",
        default=None)
    input_parser.add_outlier_rejection(default=1)
    input_parser.add_argument(
        "--sda",
        "-sda",
        action='store_true',
        help="If given, the volume is reconstructed using "
        "Scattered Data Approximation (Vercauteren et al., 2006). "
        "--alpha is considered the value for the standard deviation then. "
        "Recommended value is, e.g., --alpha 0.8")

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    filename_srr = "srr"
    dir_output_preprocessing = os.path.join(args.dir_output,
                                            "preprocessing_n4itk")
    dir_output_recon_subject_space = os.path.join(args.dir_output,
                                                  "recon_subject_space")
    dir_output_recon_template_space = os.path.join(args.dir_output,
                                                   "recon_template_space")
    dir_output_data_vs_simulatd_data = os.path.join(args.dir_output,
                                                    "data_vs_simulated_data")

    srr_subject = os.path.join(dir_output_recon_subject_space,
                               "%s_subject.nii.gz" % filename_srr)
    srr_subject_mask = ph.append_to_filename(srr_subject, "_mask")
    srr_template = os.path.join(dir_output_recon_template_space,
                                "%s_template.nii.gz" % filename_srr)
    srr_template_mask = ph.append_to_filename(srr_template, "_mask")
    trafo_template = os.path.join(dir_output_recon_template_space,
                                  "registration_transform_sitk.txt")

    if args.target_stack is None:
        target_stack = args.filenames[0]
    else:
        target_stack = args.target_stack

    if args.bias_field_correction and args.run_bias_field_correction:
        for i, f in enumerate(args.filenames):
            output = os.path.join(dir_output_preprocessing,
                                  os.path.basename(f))
            cmd_args = []
            cmd_args.append("--filename %s" % f)
            cmd_args.append("--filename-mask %s" % args.filenames_masks[i])
            cmd_args.append("--output %s" % output)
            # cmd_args.append("--verbose %d" % args.verbose)
            cmd = "niftymic_correct_bias_field %s" % (" ").join(cmd_args)
            time_start_bias = ph.start_timing()
            exit_code = ph.execute_command(cmd)
            if exit_code != 0:
                raise RuntimeError("Bias field correction failed")
        elapsed_time_bias = ph.stop_timing(time_start_bias)
        filenames = [
            os.path.join(dir_output_preprocessing, os.path.basename(f))
            for f in args.filenames
        ]
    elif args.bias_field_correction and not args.run_bias_field_correction:
        elapsed_time_bias = ph.get_zero_time()
        filenames = [
            os.path.join(dir_output_preprocessing, os.path.basename(f))
            for f in args.filenames
        ]
    else:
        elapsed_time_bias = ph.get_zero_time()
        filenames = args.filenames

    if args.run_recon_subject_space:
        target_stack_index = args.filenames.index(target_stack)

        cmd_args = []
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        if args.filenames_masks is not None:
            cmd_args.append("--filenames-masks %s" %
                            (" ").join(args.filenames_masks))
        cmd_args.append("--multiresolution %d" % args.multiresolution)
        cmd_args.append("--target-stack-index %d" % target_stack_index)
        cmd_args.append("--output %s" % srr_subject)
        cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
        cmd_args.append("--intensity-correction %d" %
                        args.intensity_correction)
        cmd_args.append("--alpha %s" % args.alpha)
        cmd_args.append("--iter-max %d" % args.iter_max)
        cmd_args.append("--two-step-cycles %d" % args.two_step_cycles)
        cmd_args.append("--outlier-rejection %d" % args.outlier_rejection)
        cmd_args.append("--verbose %d" % args.verbose)
        if args.isotropic_resolution is not None:
            cmd_args.append("--isotropic-resolution %f" %
                            args.isotropic_resolution)
        if args.reference is not None:
            cmd_args.append("--reference %s" % args.reference)
        if args.reference_mask is not None:
            cmd_args.append("--reference-mask %s" % args.reference_mask)
        if args.sda:
            cmd_args.append("--sda")
        cmd = "niftymic_reconstruct_volume %s" % (" ").join(cmd_args)
        time_start_volrec = ph.start_timing()
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Reconstruction in subject space failed")

        # Compute SRR mask in subject space
        # (Approximated using SDA within reconstruct_volume)
        if 0:
            dir_motion_correction = os.path.join(
                dir_output_recon_subject_space, "motion_correction")
            cmd_args = ["niftymic_reconstruct_volume_from_slices"]
            cmd_args.append("--filenames %s" % " ".join(args.filenames_masks))
            cmd_args.append("--dir-input-mc %s" % dir_motion_correction)
            cmd_args.append("--output %s" % srr_subject_mask)
            cmd_args.append("--reconstruction-space %s" % srr_subject)
            cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
            cmd_args.append("--mask")
            if args.sda:
                cmd_args.append("--sda")
                cmd_args.append("--alpha 1")
            else:
                cmd_args.append("--alpha 0.1")
                cmd_args.append("--iter-max 5")
            cmd = (" ").join(cmd_args)
            ph.execute_command(cmd)

        elapsed_time_volrec = ph.stop_timing(time_start_volrec)
    else:
        elapsed_time_volrec = ph.get_zero_time()

    if args.run_recon_template_space:

        if args.gestational_age is None:
            template_stack_estimator = \
                tse.TemplateStackEstimator.from_mask(srr_subject_mask)
            gestational_age = template_stack_estimator.get_estimated_gw()
            ph.print_info("Estimated gestational age: %d" % gestational_age)
        else:
            gestational_age = args.gestational_age

        template = os.path.join(DIR_TEMPLATES,
                                "STA%d.nii.gz" % gestational_age)
        template_mask = os.path.join(DIR_TEMPLATES,
                                     "STA%d_mask.nii.gz" % gestational_age)

        cmd_args = []
        cmd_args.append("--fixed %s" % template)
        cmd_args.append("--moving %s" % srr_subject)
        cmd_args.append("--fixed-mask %s" % template_mask)
        cmd_args.append("--moving-mask %s" % srr_subject_mask)
        cmd_args.append(
            "--dir-input-mc %s" %
            os.path.join(dir_output_recon_subject_space, "motion_correction"))
        cmd_args.append("--output %s" % trafo_template)
        cmd_args.append("--verbose %s" % args.verbose)
        if args.initial_transform is not None:
            cmd_args.append("--initial-transform %s" % args.initial_transform)
            cmd_args.append("--use-flirt 0")
            cmd_args.append("--test-ap-flip 0")
        cmd = "niftymic_register_image %s" % (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Registration to template space failed")

        # reconstruct volume in template space
        # pattern = "[a-zA-Z0-9_.]+(ResamplingToTemplateSpace.nii.gz)"
        # p = re.compile(pattern)
        # reconstruction_space = [
        #     os.path.join(dir_output_recon_template_space, p.match(f).group(0))
        #     for f in os.listdir(dir_output_recon_template_space)
        #     if p.match(f)][0]

        dir_input_mc = os.path.join(dir_output_recon_template_space,
                                    "motion_correction")
        cmd_args = ["niftymic_reconstruct_volume_from_slices"]
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        cmd_args.append("--dir-input-mc %s" % dir_input_mc)
        cmd_args.append("--output %s" % srr_template)
        cmd_args.append("--reconstruction-space %s" % template)
        cmd_args.append("--iter-max %d" % args.iter_max)
        cmd_args.append("--alpha %s" % args.alpha)
        cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
        cmd_args.append("--verbose %s" % args.verbose)
        if args.sda:
            cmd_args.append("--sda")

        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Reconstruction in template space failed")

        # Compute SRR mask in template space
        if 1:
            dir_motion_correction = os.path.join(
                dir_output_recon_template_space, "motion_correction")
            cmd_args = ["niftymic_reconstruct_volume_from_slices"]
            cmd_args.append("--filenames %s" % " ".join(args.filenames_masks))
            cmd_args.append("--dir-input-mc %s" % dir_motion_correction)
            cmd_args.append("--output %s" % srr_template_mask)
            cmd_args.append("--reconstruction-space %s" % srr_template)
            cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
            cmd_args.append("--mask")
            if args.sda:
                cmd_args.append("--sda")
                cmd_args.append("--alpha 1")
            else:
                cmd_args.append("--alpha 0.1")
                cmd_args.append("--iter-max 5")
            cmd = (" ").join(cmd_args)
            ph.execute_command(cmd)

        # Copy SRR to output directory
        output = "%sSRR_Stacks%d.nii.gz" % (args.prefix_output,
                                            len(args.filenames))
        path_to_output = os.path.join(args.dir_output, output)
        cmd = "cp -p %s %s" % (srr_template, path_to_output)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Copy of SRR to output directory failed")

        # Multiply template mask with reconstruction
        cmd_args = ["niftymic_multiply"]
        fnames = [
            srr_template,
            srr_template_mask,
        ]
        output_masked = "Masked_%s" % output
        path_to_output_masked = os.path.join(args.dir_output, output_masked)
        cmd_args.append("--filenames %s" % " ".join(fnames))
        cmd_args.append("--output %s" % path_to_output_masked)
        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("SRR brain masking failed")

    else:
        elapsed_time_template = ph.get_zero_time()

    if args.run_data_vs_simulated_data:

        dir_input_mc = os.path.join(dir_output_recon_template_space,
                                    "motion_correction")

        # Get simulated/projected slices
        cmd_args = []
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        if args.filenames_masks is not None:
            cmd_args.append("--filenames-masks %s" %
                            (" ").join(args.filenames_masks))
        cmd_args.append("--dir-input-mc %s" % dir_input_mc)
        cmd_args.append("--dir-output %s" % dir_output_data_vs_simulatd_data)
        cmd_args.append("--reconstruction %s" % srr_template)
        cmd_args.append("--copy-data 1")
        cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
        # cmd_args.append("--verbose %s" % args.verbose)
        exe = os.path.abspath(simulate_stacks_from_reconstruction.__file__)
        cmd = "python %s %s" % (exe, (" ").join(cmd_args))
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("SRR slice projections failed")

        filenames_simulated = [
            os.path.join(dir_output_data_vs_simulatd_data, os.path.basename(f))
            for f in filenames
        ]

        dir_output_evaluation = os.path.join(dir_output_data_vs_simulatd_data,
                                             "evaluation")
        dir_output_figures = os.path.join(dir_output_data_vs_simulatd_data,
                                          "figures")
        dir_output_side_by_side = os.path.join(dir_output_figures,
                                               "side-by-side")

        # Evaluate slice similarities to ground truth
        cmd_args = []
        cmd_args.append("--filenames %s" % (" ").join(filenames_simulated))
        if args.filenames_masks is not None:
            cmd_args.append("--filenames-masks %s" %
                            (" ").join(args.filenames_masks))
        cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
        cmd_args.append("--measures NCC SSIM")
        cmd_args.append("--dir-output %s" % dir_output_evaluation)
        exe = os.path.abspath(evaluate_simulated_stack_similarity.__file__)
        cmd = "python %s %s" % (exe, (" ").join(cmd_args))
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Evaluation of slice similarities failed")

        # Generate figures showing the quantitative comparison
        cmd_args = []
        cmd_args.append("--dir-input %s" % dir_output_evaluation)
        cmd_args.append("--dir-output %s" % dir_output_figures)
        exe = os.path.abspath(
            show_evaluated_simulated_stack_similarity.__file__)
        cmd = "python %s %s" % (exe, (" ").join(cmd_args))
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            ph.print_warning("Visualization of slice similarities failed")

        # Generate pdfs showing all the side-by-side comparisons
        cmd_args = []
        cmd_args.append("--filenames %s" % (" ").join(filenames_simulated))
        cmd_args.append("--dir-output %s" % dir_output_side_by_side)
        exe = os.path.abspath(
            export_side_by_side_simulated_vs_original_slice_comparison.__file__
        )
        cmd = "python %s %s" % (exe, (" ").join(cmd_args))
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Generation of PDF overview failed")

    ph.print_title("Summary")
    print("Computational Time for Bias Field Correction: %s" %
          elapsed_time_bias)
    print("Computational Time for Volumetric Reconstruction: %s" %
          elapsed_time_volrec)
    print("Computational Time for Pipeline: %s" % ph.stop_timing(time_start))

    return 0
Beispiel #18
0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Register an obtained reconstruction (moving) "
        "to a template image/space (fixed) using rigid registration. "
        "The resulting registration can optionally be applied to previously "
        "obtained motion correction slice transforms so that a volumetric "
        "reconstruction is possible in the (standard anatomical) space "
        "defined by the fixed.", )
    input_parser.add_fixed(required=True)
    input_parser.add_moving(required=True)
    input_parser.add_output(help="Path to registration transform (.txt)",
                            required=True)
    input_parser.add_fixed_mask(required=False)
    input_parser.add_moving_mask(required=False)
    input_parser.add_option(
        option_string="--initial-transform",
        type=str,
        help="Path to initial transform. "
        "If not provided, registration will be initialized based on "
        "rigid alignment of eigenbasis of the fixed/moving image masks "
        "using principal component analysis",
        default=None)
    input_parser.add_v2v_method(
        option_string="--method",
        help="Registration method used for the registration.",
        default="RegAladin",
    )
    input_parser.add_argument(
        "--refine-pca",
        "-refine-pca",
        action='store_true',
        help="If given, PCA-based initializations will be refined using "
        "RegAladin registrations.")
    input_parser.add_dir_input_mc()
    input_parser.add_verbose(default=0)
    input_parser.add_log_config(default=1)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    if not args.output.endswith(".txt"):
        raise IOError("output transformation path must end in '.txt'")

    dir_output = os.path.dirname(args.output)
    ph.create_directory(dir_output)

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")
    fixed = st.Stack.from_filename(file_path=args.fixed,
                                   file_path_mask=args.fixed_mask,
                                   extract_slices=False)
    moving = st.Stack.from_filename(file_path=args.moving,
                                    file_path_mask=args.moving_mask,
                                    extract_slices=False)

    path_to_tmp_output = os.path.join(
        DIR_TMP, ph.append_to_filename(os.path.basename(args.moving),
                                       "_warped"))

    # ---------------------------- Initialization ----------------------------
    if args.initial_transform is None:
        ph.print_title("Estimate initial transform using PCA")

        if args.moving_mask is None or args.fixed_mask is None:
            ph.print_warning("Fixed and moving masks are strongly recommended")
        transform_initializer = tinit.TransformInitializer(
            fixed=fixed,
            moving=moving,
            similarity_measure="NMI",
            refine_pca_initializations=args.refine_pca,
        )
        transform_initializer.run()
        transform_init_sitk = transform_initializer.get_transform_sitk()
    else:
        transform_init_sitk = sitkh.read_transform_sitk(args.initial_transform)
    sitk.WriteTransform(transform_init_sitk, args.output)

    # -------------------Register Reconstruction to Template-------------------
    ph.print_title("Registration")

    if args.method == "RegAladin":

        path_to_transform_regaladin = os.path.join(DIR_TMP,
                                                   "transform_regaladin.txt")

        # Convert SimpleITK to RegAladin transform
        cmd = "simplereg_transform -sitk2nreg %s %s" % (
            args.output, path_to_transform_regaladin)
        ph.execute_command(cmd, verbose=False)

        # Run NiftyReg
        cmd_args = ["reg_aladin"]
        cmd_args.append("-ref '%s'" % args.fixed)
        cmd_args.append("-flo '%s'" % args.moving)
        cmd_args.append("-res '%s'" % path_to_tmp_output)
        cmd_args.append("-inaff '%s'" % path_to_transform_regaladin)
        cmd_args.append("-aff '%s'" % path_to_transform_regaladin)
        cmd_args.append("-rigOnly")
        cmd_args.append("-ln 2")  # seems to perform better for spina bifida
        cmd_args.append("-voff")
        if args.fixed_mask is not None:
            cmd_args.append("-rmask '%s'" % args.fixed_mask)

        # To avoid error "0 correspondences between blocks were found" that can
        # occur for some cases. Also, disable moving mask, as this would be ignored
        # anyway
        cmd_args.append("-noSym")
        # if args.moving_mask is not None:
        #     cmd_args.append("-fmask '%s'" % args.moving_mask)

        ph.print_info("Run Registration (RegAladin) ... ", newline=False)
        ph.execute_command(" ".join(cmd_args), verbose=False)
        print("done")

        # Convert RegAladin to SimpleITK transform
        cmd = "simplereg_transform -nreg2sitk '%s' '%s'" % (
            path_to_transform_regaladin, args.output)
        ph.execute_command(cmd, verbose=False)

    else:
        path_to_transform_flirt = os.path.join(DIR_TMP, "transform_flirt.txt")

        # Convert SimpleITK into FLIRT transform
        cmd = "simplereg_transform -sitk2flirt '%s' '%s' '%s' '%s'" % (
            args.output, args.fixed, args.moving, path_to_transform_flirt)
        ph.execute_command(cmd, verbose=False)

        # Define search angle ranges for FLIRT in all three dimensions
        search_angles = [
            "-searchr%s -%d %d" % (x, 180, 180) for x in ["x", "y", "z"]
        ]

        cmd_args = ["flirt"]
        cmd_args.append("-in '%s'" % args.moving)
        cmd_args.append("-ref '%s'" % args.fixed)
        if args.initial_transform is not None:
            cmd_args.append("-init '%s'" % path_to_transform_flirt)
        cmd_args.append("-omat '%s'" % path_to_transform_flirt)
        cmd_args.append("-out '%s'" % path_to_tmp_output)
        cmd_args.append("-dof 6")
        cmd_args.append((" ").join(search_angles))
        if args.moving_mask is not None:
            cmd_args.append("-inweight '%s'" % args.moving_mask)
        if args.fixed_mask is not None:
            cmd_args.append("-refweight '%s'" % args.fixed_mask)
        ph.print_info("Run Registration (FLIRT) ... ", newline=False)
        ph.execute_command(" ".join(cmd_args), verbose=False)
        print("done")

        # Convert FLIRT to SimpleITK transform
        cmd = "simplereg_transform -flirt2sitk '%s' '%s' '%s' '%s'" % (
            path_to_transform_flirt, args.fixed, args.moving, args.output)
        ph.execute_command(cmd, verbose=False)

    if args.dir_input_mc is not None:
        ph.print_title("Update Motion-Correction Transformations")
        transform_sitk = sitkh.read_transform_sitk(args.output, inverse=1)

        if args.dir_input_mc.endswith("/"):
            subdir_mc = args.dir_input_mc.split("/")[-2]
        else:
            subdir_mc = args.dir_input_mc.split("/")[-1]
        dir_output_mc = os.path.join(dir_output, subdir_mc)

        ph.create_directory(dir_output_mc, delete_files=True)
        pattern = REGEX_FILENAMES + "[.]tfm"
        p = re.compile(pattern)
        trafos = [t for t in os.listdir(args.dir_input_mc) if p.match(t)]
        for t in trafos:
            path_to_input_transform = os.path.join(args.dir_input_mc, t)
            path_to_output_transform = os.path.join(dir_output_mc, t)
            t_sitk = sitkh.read_transform_sitk(path_to_input_transform)
            t_sitk = sitkh.get_composite_sitk_affine_transform(
                transform_sitk, t_sitk)
            sitk.WriteTransform(t_sitk, path_to_output_transform)
        ph.print_info("%d transformations written to '%s'" %
                      (len(trafos), dir_output_mc))

    if args.verbose:
        ph.show_niftis([args.fixed, path_to_tmp_output])

    elapsed_time_total = ph.stop_timing(time_start)

    # Summary
    ph.print_title("Summary")
    print("Computational Time: %s" % (elapsed_time_total))

    return 0
def main():
    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Perform automatic brain masking using "
        "fetal_brain_seg, part of the MONAIfbs package "
        "(https://github.com/gift-surg/MONAIfbs). ",
    )
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks(required=False)
    input_parser.add_dir_output(required=False)
    input_parser.add_verbose(default=0)
    input_parser.add_log_config(default=0)
    input_parser.add_option(
        option_string="--neuroimage-legacy-seg",
        type=int,
        required=False,
        default=0,
        help="If set to 1, use the legacy method for fetal brain segmentation "
             "i.e. the two-step approach proposed in Ebner, Wang et al "
             "NeuroImage (2020)"
    )

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.neuroimage_legacy_seg:
        try:
            DIR_FETAL_BRAIN_SEG = os.environ["FETAL_BRAIN_SEG"]
        except KeyError as e:
            raise RuntimeError(
                "Environment variable FETAL_BRAIN_SEG is not specified. "
                "Specify the root directory of fetal_brain_seg "
                "(https://github.com/gift-surg/fetal_brain_seg) "
                "using "
                "'export FETAL_BRAIN_SEG=path_to_fetal_brain_seg_dir' "
                "(in bashrc).")
    else:
        try:
            import monaifbs
            DIR_FETAL_BRAIN_SEG = os.path.dirname(monaifbs.__file__)
        except ImportError as e:
            raise RuntimeError(
                "monaifbs not correctly installed. "
                "Please check its installation running "
                "pip install -e MONAIfbs/ "
            )

    print("Using executable from {}".format(DIR_FETAL_BRAIN_SEG))

    if args.filenames_masks is None and args.dir_output is None:
        raise IOError("Either --filenames-masks or --dir-output must be set")

    if args.dir_output is not None:
        args.filenames_masks = [
            os.path.join(args.dir_output, os.path.basename(f))
            for f in args.filenames
        ]

    if len(args.filenames) != len(args.filenames_masks):
        raise IOError("Number of filenames and filenames-masks must match")

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    cd_fetal_brain_seg = "cd %s" % DIR_FETAL_BRAIN_SEG

    for f, m in zip(args.filenames, args.filenames_masks):

        if not ph.file_exists(f):
            raise IOError("File '%s' does not exist" % f)

        # use absolute path for input image
        f = os.path.abspath(f)

        # use absolute path for output image
        dir_output = os.path.dirname(m)
        if not os.path.isabs(dir_output):
            dir_output = os.path.realpath(
                os.path.join(os.getcwd(), dir_output))
            m = os.path.join(dir_output, os.path.basename(m))

        ph.create_directory(dir_output)

        # Change to root directory of fetal_brain_seg
        cmds = [cd_fetal_brain_seg]

        # Run masking independently (Takes longer but ensures that it does
        # not terminate because of provided 'non-brain images')
        cmd_args = ["python fetal_brain_seg.py"]
        cmd_args.append("--input_names '%s'" % f)
        cmd_args.append("--segment_output_names '%s'" % m)
        cmds.append(" ".join(cmd_args))

        # Execute both steps
        cmd = " && ".join(cmds)
        flag = ph.execute_command(cmd)

        if flag != 0:
            ph.print_warning(
                "Error using fetal_brain_seg. \n"
                "Execute '%s' for further investigation" %
                cmd)

        ph.print_info("Fetal brain segmentation written to '%s'" % m)

        if args.verbose:
            ph.show_nifti(f, segmentation=m)

    elapsed_time_total = ph.stop_timing(time_start)

    ph.print_title("Summary")
    exe_file_info = os.path.basename(os.path.abspath(__file__)).split(".")[0]
    print("%s | Computational Time: %s" % (exe_file_info, elapsed_time_total))

    return 0
Beispiel #20
0
def main():

    # Set print options
    np.set_printoptions(precision=3)
    pd.set_option('display.width', 1000)

    input_parser = InputArgparser(description=".", )
    input_parser.add_filenames(required=True)
    input_parser.add_reference(required=True)
    input_parser.add_reference_mask()
    input_parser.add_dir_output(required=False)
    input_parser.add_measures(
        default=["PSNR", "RMSE", "MAE", "SSIM", "NCC", "NMI"])
    input_parser.add_verbose(default=0)
    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    ph.print_title("Image similarity")
    data_reader = dr.MultipleImagesReader(args.filenames)
    data_reader.read_data()
    stacks = data_reader.get_data()

    reference = st.Stack.from_filename(args.reference, args.reference_mask)

    for stack in stacks:
        try:
            stack.sitk - reference.sitk
        except RuntimeError as e:
            raise IOError(
                "All provided images must be at the same image space")

    x_ref = sitk.GetArrayFromImage(reference.sitk)

    if args.reference_mask is None:
        indices = np.where(x_ref != np.inf)
    else:
        x_ref_mask = sitk.GetArrayFromImage(reference.sitk_mask)
        indices = np.where(x_ref_mask > 0)

    measures_dic = {
        m: lambda x, m=m: SimilarityMeasures.similarity_measures[m]
        (x[indices], x_ref[indices])
        # SimilarityMeasures.similarity_measures[m](x, x_ref)
        for m in args.measures
    }

    observer = obs.Observer()
    observer.set_measures(measures_dic)
    for stack in stacks:
        nda = sitk.GetArrayFromImage(stack.sitk)
        observer.add_x(nda)

    if args.verbose:
        stacks_comparison = [s for s in stacks]
        stacks_comparison.insert(0, reference)
        sitkh.show_stacks(
            stacks_comparison,
            segmentation=reference,
        )

    observer.compute_measures()
    measures = observer.get_measures()

    # Store information in array
    error = np.zeros((len(stacks), len(measures)))
    cols = measures
    rows = []
    for i_stack, stack in enumerate(stacks):
        error[i_stack, :] = np.array([measures[m][i_stack] for m in measures])
        rows.append(stack.get_filename())

    header = "# Ref: %s, Ref-Mask: %d, %s \n" % (
        reference.get_filename(),
        args.reference_mask is None,
        ph.get_time_stamp(),
    )
    header += "# %s\n" % ("\t").join(measures)

    path_to_file_filenames = os.path.join(args.dir_output, "filenames.txt")
    path_to_file_similarities = os.path.join(args.dir_output,
                                             "similarities.txt")

    # Write to files
    ph.write_to_file(path_to_file_similarities, header)
    ph.write_array_to_file(path_to_file_similarities, error, verbose=False)
    text = header
    text += "%s\n" % "\n".join(rows)
    ph.write_to_file(path_to_file_filenames, text)

    # Print to screen
    ph.print_subtitle("Computed Similarities")
    df = pd.DataFrame(error, rows, cols)
    print(df)

    return 0
Beispiel #21
0
    def _run(self):

        ph.print_title("Slice-to-Volume Registration")

        self._registration_method.set_moving(self._reference)

        for i, stack in enumerate(self._stacks):
            slices = stack.get_slices()

            transforms_sitk = {}

            for j, slice_j in enumerate(slices):

                txt = "%sSlice-to-Volume Registration -- " \
                    "Stack %d/%d (%s) -- Slice %d/%d" % (
                        self._print_prefix,
                        i + 1, len(self._stacks), stack.get_filename(),
                        j + 1, len(slices))
                if self._verbose:
                    ph.print_subtitle(txt)
                else:
                    ph.print_info(txt)

                self._registration_method.set_fixed(slice_j)
                self._registration_method.run()

                # Store information on registration transform
                transform_sitk = \
                    self._registration_method.get_registration_transform_sitk()
                transforms_sitk[slice_j.get_slice_number()] = transform_sitk

            # Avoid slice misregistrations
            if self._s2v_smoothing is not None:
                # import os
                # for slice_number in transforms_sitk.keys():
                #     path_to_file = os.path.join(
                #         "/tmp/fetal_brain", "%s_slice%d.tfm" % (
                #             stack.get_filename(), slice_number))
                #     sitk.WriteTransform(
                #         transforms_sitk[slice_number], path_to_file)
                ph.print_subtitle("Robust slice motion estimation "
                                  "(GP smoothing = %g, interleave = %d)" %
                                  (self._s2v_smoothing, self._interleave))
                robust_motion_estimator = rme.RobustMotionEstimator(
                    transforms_sitk=transforms_sitk,
                    interleave=self._interleave)
                robust_motion_estimator.run(self._s2v_smoothing)
                transforms_sitk = \
                    robust_motion_estimator.get_robust_transforms_sitk()

                # Update position of slice
                for slice in slices:
                    slice_number = slice.get_slice_number()
                    slice.update_motion_correction(
                        transforms_sitk[slice_number])

                # Run s2v-reg again
                for j, slice_j in enumerate(slices):
                    txt = "%sSlice-to-Volume Registration -- " \
                        "Stack %d/%d -- Slice %d/%d (after GP init)" % (
                            self._print_prefix,
                            i + 1, len(self._stacks),
                            j + 1, len(slices))
                    if self._verbose:
                        ph.print_subtitle(txt)
                    else:
                        ph.print_info(txt)

                    self._registration_method.set_fixed(slice_j)
                    self._registration_method.run()

                    # Store information on registration transform
                    transform_sitk = \
                        self._registration_method.get_registration_transform_sitk()
                    transforms_sitk[
                        slice_j.get_slice_number()] = transform_sitk

                # Export figures
                # title = "%s_Stack%d%s" % (
                #     self._print_prefix, i, stack.get_filename())
                # title = ph.replace_string_for_print(title)
                # robust_motion_estimator.show_estimated_transform_parameters(
                #     dir_output="/tmp/fetal_brain/figs", title=title)

            # dir_output = "/tmp/fetal/figs"
            # motion_evaluator = me.MotionEvaluator(transforms_sitk)
            # motion_evaluator.run()
            # motion_evaluator.display(dir_output=dir_output, title=title)
            # motion_evaluator.show(dir_output=dir_output, title=title)

            # Update position of slice
            for slice in slices:
                slice_number = slice.get_slice_number()
                slice.update_motion_correction(transforms_sitk[slice_number])
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Run reconstruction pipeline including "
        "(i) bias field correction, "
        "(ii) volumetric reconstruction in subject space, "
        "(iii) volumetric reconstruction in template space, "
        "and (iv) some diagnostics to assess the obtained reconstruction.", )
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks(required=True)
    input_parser.add_target_stack(required=False)
    input_parser.add_suffix_mask(default="")
    input_parser.add_dir_output(required=True)
    input_parser.add_alpha(default=0.01)
    input_parser.add_verbose(default=0)
    input_parser.add_gestational_age(required=False)
    input_parser.add_prefix_output(default="")
    input_parser.add_search_angle(default=180)
    input_parser.add_multiresolution(default=0)
    input_parser.add_log_config(default=1)
    input_parser.add_isotropic_resolution()
    input_parser.add_reference()
    input_parser.add_reference_mask()
    input_parser.add_bias_field_correction(default=1)
    input_parser.add_intensity_correction(default=1)
    input_parser.add_iter_max(default=10)
    input_parser.add_two_step_cycles(default=3)
    input_parser.add_slice_thicknesses(default=None)
    input_parser.add_option(
        option_string="--run-bias-field-correction",
        type=int,
        help="Turn on/off bias field correction. "
        "If off, it is assumed that this step was already performed "
        "if --bias-field-correction is active.",
        default=1)
    input_parser.add_option(
        option_string="--run-recon-subject-space",
        type=int,
        help="Turn on/off reconstruction in subject space. "
        "If off, it is assumed that this step was already performed.",
        default=1)
    input_parser.add_option(
        option_string="--run-recon-template-space",
        type=int,
        help="Turn on/off reconstruction in template space. "
        "If off, it is assumed that this step was already performed.",
        default=1)
    input_parser.add_option(
        option_string="--run-diagnostics",
        type=int,
        help="Turn on/off diagnostics of the obtained volumetric "
        "reconstruction. ",
        default=0)
    input_parser.add_option(
        option_string="--initial-transform",
        type=str,
        help="Set initial transform to be used for register_image.",
        default=None)
    input_parser.add_outlier_rejection(default=1)
    input_parser.add_threshold_first(default=0.5)
    input_parser.add_threshold(default=0.8)
    input_parser.add_argument(
        "--sda",
        "-sda",
        action='store_true',
        help="If given, the volume is reconstructed using "
        "Scattered Data Approximation (Vercauteren et al., 2006). "
        "--alpha is considered the value for the standard deviation then. "
        "Recommended value is, e.g., --alpha 0.8")
    input_parser.add_argument(
        "--v2v-robust",
        "-v2v-robust",
        action='store_true',
        help="If given, a more robust volume-to-volume registration step is "
        "performed, i.e. four rigid registrations are performed using four "
        "rigid transform initializations based on "
        "principal component alignment of associated masks.")
    input_parser.add_interleave(default=3)
    input_parser.add_argument(
        "--s2v-hierarchical",
        "-s2v-hierarchical",
        action='store_true',
        help="If given, a hierarchical approach for the first slice-to-volume "
        "registration cycle is used, i.e. sub-packages defined by the "
        "specified interleave (--interleave) are registered until each "
        "slice is registered independently.")

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    filename_srr = "srr"
    dir_output_preprocessing = os.path.join(args.dir_output,
                                            "preprocessing_n4itk")
    dir_output_recon_subject_space = os.path.join(args.dir_output,
                                                  "recon_subject_space")
    dir_output_recon_template_space = os.path.join(args.dir_output,
                                                   "recon_template_space")
    dir_output_diagnostics = os.path.join(args.dir_output, "diagnostics")

    srr_subject = os.path.join(dir_output_recon_subject_space,
                               "%s_subject.nii.gz" % filename_srr)
    srr_subject_mask = ph.append_to_filename(srr_subject, "_mask")
    srr_template = os.path.join(dir_output_recon_template_space,
                                "%s_template.nii.gz" % filename_srr)
    srr_template_mask = ph.append_to_filename(srr_template, "_mask")
    trafo_template = os.path.join(dir_output_recon_template_space,
                                  "registration_transform_sitk.txt")
    srr_slice_coverage = os.path.join(
        dir_output_diagnostics,
        "%s_template_slicecoverage.nii.gz" % filename_srr)

    if args.bias_field_correction and args.run_bias_field_correction:
        for i, f in enumerate(args.filenames):
            output = os.path.join(dir_output_preprocessing,
                                  os.path.basename(f))
            cmd_args = []
            cmd_args.append("--filename '%s'" % f)
            cmd_args.append("--filename-mask '%s'" % args.filenames_masks[i])
            cmd_args.append("--output '%s'" % output)
            # cmd_args.append("--verbose %d" % args.verbose)
            cmd_args.append("--log-config %d" % args.log_config)
            cmd = "niftymic_correct_bias_field %s" % (" ").join(cmd_args)
            time_start_bias = ph.start_timing()
            exit_code = ph.execute_command(cmd)
            if exit_code != 0:
                raise RuntimeError("Bias field correction failed")
        elapsed_time_bias = ph.stop_timing(time_start_bias)
        filenames = [
            os.path.join(dir_output_preprocessing, os.path.basename(f))
            for f in args.filenames
        ]
    elif args.bias_field_correction and not args.run_bias_field_correction:
        elapsed_time_bias = ph.get_zero_time()
        filenames = [
            os.path.join(dir_output_preprocessing, os.path.basename(f))
            for f in args.filenames
        ]
    else:
        elapsed_time_bias = ph.get_zero_time()
        filenames = args.filenames

    # Specify target stack for intensity correction and reconstruction space
    if args.target_stack is None:
        target_stack = filenames[0]
    else:
        try:
            target_stack_index = args.filenames.index(args.target_stack)
        except ValueError as e:
            raise ValueError(
                "--target-stack must correspond to an image as provided by "
                "--filenames")
        target_stack = filenames[target_stack_index]

    # Add single quotes around individual filenames to account for whitespaces
    filenames = ["'" + f + "'" for f in filenames]
    filenames_masks = ["'" + f + "'" for f in args.filenames_masks]

    if args.run_recon_subject_space:

        cmd_args = ["niftymic_reconstruct_volume"]
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        cmd_args.append("--filenames-masks %s" % (" ").join(filenames_masks))
        cmd_args.append("--multiresolution %d" % args.multiresolution)
        cmd_args.append("--target-stack '%s'" % target_stack)
        cmd_args.append("--output '%s'" % srr_subject)
        cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
        cmd_args.append("--intensity-correction %d" %
                        args.intensity_correction)
        cmd_args.append("--alpha %s" % args.alpha)
        cmd_args.append("--iter-max %d" % args.iter_max)
        cmd_args.append("--two-step-cycles %d" % args.two_step_cycles)
        cmd_args.append("--outlier-rejection %d" % args.outlier_rejection)
        cmd_args.append("--threshold-first %f" % args.threshold_first)
        cmd_args.append("--threshold %f" % args.threshold)
        if args.slice_thicknesses is not None:
            cmd_args.append("--slice-thicknesses %s" %
                            " ".join(map(str, args.slice_thicknesses)))
        cmd_args.append("--verbose %d" % args.verbose)
        cmd_args.append("--log-config %d" % args.log_config)
        if args.isotropic_resolution is not None:
            cmd_args.append("--isotropic-resolution %f" %
                            args.isotropic_resolution)
        if args.reference is not None:
            cmd_args.append("--reference %s" % args.reference)
        if args.reference_mask is not None:
            cmd_args.append("--reference-mask %s" % args.reference_mask)
        if args.sda:
            cmd_args.append("--sda")
        if args.v2v_robust:
            cmd_args.append("--v2v-robust")
        if args.s2v_hierarchical:
            cmd_args.append("--s2v-hierarchical")

        cmd = (" ").join(cmd_args)
        time_start_volrec = ph.start_timing()
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Reconstruction in subject space failed")

        # Compute SRR mask in subject space
        # (Approximated using SDA within reconstruct_volume)
        if 0:
            dir_motion_correction = os.path.join(
                dir_output_recon_subject_space, "motion_correction")
            cmd_args = ["niftymic_reconstruct_volume_from_slices"]
            cmd_args.append("--filenames %s" % " ".join(filenames_masks))
            cmd_args.append("--dir-input-mc '%s'" % dir_motion_correction)
            cmd_args.append("--output '%s'" % srr_subject_mask)
            cmd_args.append("--reconstruction-space '%s'" % srr_subject)
            cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
            cmd_args.append("--mask")
            cmd_args.append("--log-config %d" % args.log_config)
            if args.slice_thicknesses is not None:
                cmd_args.append("--slice-thicknesses %s" %
                                " ".join(map(str, args.slice_thicknesses)))
            if args.sda:
                cmd_args.append("--sda")
                cmd_args.append("--alpha 1")
            else:
                cmd_args.append("--alpha 0.1")
                cmd_args.append("--iter-max 5")
            cmd = (" ").join(cmd_args)
            ph.execute_command(cmd)

        elapsed_time_volrec = ph.stop_timing(time_start_volrec)
    else:
        elapsed_time_volrec = ph.get_zero_time()

    if args.run_recon_template_space:

        if args.gestational_age is None:
            template_stack_estimator = \
                tse.TemplateStackEstimator.from_mask(srr_subject_mask)
            gestational_age = template_stack_estimator.get_estimated_gw()
            ph.print_info("Estimated gestational age: %d" % gestational_age)
        else:
            gestational_age = args.gestational_age

        template = os.path.join(DIR_TEMPLATES,
                                "STA%d.nii.gz" % gestational_age)
        template_mask = os.path.join(DIR_TEMPLATES,
                                     "STA%d_mask.nii.gz" % gestational_age)

        # Register SRR to template space
        cmd_args = ["niftymic_register_image"]
        cmd_args.append("--fixed '%s'" % template)
        cmd_args.append("--moving '%s'" % srr_subject)
        cmd_args.append("--fixed-mask '%s'" % template_mask)
        cmd_args.append("--moving-mask '%s'" % srr_subject_mask)
        cmd_args.append(
            "--dir-input-mc '%s'" %
            os.path.join(dir_output_recon_subject_space, "motion_correction"))
        cmd_args.append("--output '%s'" % trafo_template)
        cmd_args.append("--verbose %s" % args.verbose)
        cmd_args.append("--log-config %d" % args.log_config)
        cmd_args.append("--refine-pca")
        if args.initial_transform is not None:
            cmd_args.append("--initial-transform '%s'" %
                            args.initial_transform)
        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Registration to template space failed")

        # Compute SRR in template space
        dir_input_mc = os.path.join(dir_output_recon_template_space,
                                    "motion_correction")
        cmd_args = ["niftymic_reconstruct_volume_from_slices"]
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        cmd_args.append("--filenames-masks %s" % (" ").join(filenames_masks))
        cmd_args.append("--dir-input-mc '%s'" % dir_input_mc)
        cmd_args.append("--output '%s'" % srr_template)
        cmd_args.append("--reconstruction-space '%s'" % template)
        cmd_args.append("--target-stack '%s'" % target_stack)
        cmd_args.append("--iter-max %d" % args.iter_max)
        cmd_args.append("--alpha %s" % args.alpha)
        cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
        cmd_args.append("--verbose %s" % args.verbose)
        cmd_args.append("--log-config %d" % args.log_config)
        if args.slice_thicknesses is not None:
            cmd_args.append("--slice-thicknesses %s" %
                            " ".join(map(str, args.slice_thicknesses)))
        if args.sda:
            cmd_args.append("--sda")

        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Reconstruction in template space failed")

        # Compute SRR mask in template space
        if 1:
            dir_motion_correction = os.path.join(
                dir_output_recon_template_space, "motion_correction")
            cmd_args = ["niftymic_reconstruct_volume_from_slices"]
            cmd_args.append("--filenames %s" % " ".join(filenames_masks))
            cmd_args.append("--dir-input-mc '%s'" % dir_motion_correction)
            cmd_args.append("--output '%s'" % srr_template_mask)
            cmd_args.append("--reconstruction-space '%s'" % srr_template)
            cmd_args.append("--suffix-mask '%s'" % args.suffix_mask)
            cmd_args.append("--log-config %d" % args.log_config)
            cmd_args.append("--mask")
            if args.slice_thicknesses is not None:
                cmd_args.append("--slice-thicknesses %s" %
                                " ".join(map(str, args.slice_thicknesses)))
            if args.sda:
                cmd_args.append("--sda")
                cmd_args.append("--alpha 1")
            else:
                cmd_args.append("--alpha 0.1")
                cmd_args.append("--iter-max 5")
            cmd = (" ").join(cmd_args)
            ph.execute_command(cmd)

        # Copy SRR to output directory
        if 0:
            output = "%sSRR_Stacks%d.nii.gz" % (args.prefix_output,
                                                len(args.filenames))
            path_to_output = os.path.join(args.dir_output, output)
            cmd = "cp -p '%s' '%s'" % (srr_template, path_to_output)
            exit_code = ph.execute_command(cmd)
            if exit_code != 0:
                raise RuntimeError("Copy of SRR to output directory failed")

        # Multiply template mask with reconstruction
        if 0:
            cmd_args = ["niftymic_multiply"]
            fnames = [
                srr_template,
                srr_template_mask,
            ]
            output_masked = "Masked_%s" % output
            path_to_output_masked = os.path.join(args.dir_output,
                                                 output_masked)
            cmd_args.append("--filenames %s" % " ".join(fnames))
            cmd_args.append("--output '%s'" % path_to_output_masked)
            cmd = (" ").join(cmd_args)
            exit_code = ph.execute_command(cmd)
            if exit_code != 0:
                raise RuntimeError("SRR brain masking failed")

    else:
        elapsed_time_template = ph.get_zero_time()

    if args.run_diagnostics:

        dir_input_mc = os.path.join(dir_output_recon_template_space,
                                    "motion_correction")
        dir_output_orig_vs_proj = os.path.join(dir_output_diagnostics,
                                               "original_vs_projected")
        dir_output_selfsimilarity = os.path.join(dir_output_diagnostics,
                                                 "selfsimilarity")
        dir_output_orig_vs_proj_pdf = os.path.join(dir_output_orig_vs_proj,
                                                   "pdf")

        # Show slice coverage over reconstruction space
        exe = os.path.abspath(show_slice_coverage.__file__)
        cmd_args = ["python %s" % exe]
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        cmd_args.append("--dir-input-mc '%s'" % dir_input_mc)
        cmd_args.append("--reconstruction-space '%s'" % srr_template)
        cmd_args.append("--output '%s'" % srr_slice_coverage)
        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Slice coverage visualization failed")

        # Get simulated/projected slices
        exe = os.path.abspath(simulate_stacks_from_reconstruction.__file__)
        cmd_args = ["python %s" % exe]
        cmd_args.append("--filenames %s" % (" ").join(filenames))
        if args.filenames_masks is not None:
            cmd_args.append("--filenames-masks %s" %
                            (" ").join(filenames_masks))
        cmd_args.append("--dir-input-mc '%s'" % dir_input_mc)
        cmd_args.append("--dir-output '%s'" % dir_output_orig_vs_proj)
        cmd_args.append("--reconstruction '%s'" % srr_template)
        cmd_args.append("--copy-data 1")
        if args.slice_thicknesses is not None:
            cmd_args.append("--slice-thicknesses %s" %
                            " ".join(map(str, args.slice_thicknesses)))
        # cmd_args.append("--verbose %s" % args.verbose)
        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("SRR slice projections failed")

        filenames_simulated = [
            "'%s" % os.path.join(dir_output_orig_vs_proj, os.path.basename(f))
            for f in filenames
        ]

        # Evaluate slice similarities to ground truth
        exe = os.path.abspath(evaluate_simulated_stack_similarity.__file__)
        cmd_args = ["python %s" % exe]
        cmd_args.append("--filenames %s" % (" ").join(filenames_simulated))
        if args.filenames_masks is not None:
            cmd_args.append("--filenames-masks %s" %
                            (" ").join(filenames_masks))
        cmd_args.append("--measures NCC SSIM")
        cmd_args.append("--dir-output '%s'" % dir_output_selfsimilarity)
        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            raise RuntimeError("Evaluation of stack similarities failed")

        # Generate figures showing the quantitative comparison
        exe = os.path.abspath(
            show_evaluated_simulated_stack_similarity.__file__)
        cmd_args = ["python %s" % exe]
        cmd_args.append("--dir-input '%s'" % dir_output_selfsimilarity)
        cmd_args.append("--dir-output '%s'" % dir_output_selfsimilarity)
        cmd = (" ").join(cmd_args)
        exit_code = ph.execute_command(cmd)
        if exit_code != 0:
            ph.print_warning("Visualization of stack similarities failed")

        # Generate pdfs showing all the side-by-side comparisons
        if 0:
            exe = os.path.abspath(
                export_side_by_side_simulated_vs_original_slice_comparison.
                __file__)
            cmd_args = ["python %s" % exe]
            cmd_args.append("--filenames %s" % (" ").join(filenames_simulated))
            cmd_args.append("--dir-output '%s'" % dir_output_orig_vs_proj_pdf)
            cmd = "python %s %s" % (exe, (" ").join(cmd_args))
            cmd = (" ").join(cmd_args)
            exit_code = ph.execute_command(cmd)
            if exit_code != 0:
                raise RuntimeError("Generation of PDF overview failed")

    ph.print_title("Summary")
    print("Computational Time for Bias Field Correction: %s" %
          elapsed_time_bias)
    print("Computational Time for Volumetric Reconstruction: %s" %
          elapsed_time_volrec)
    print("Computational Time for Pipeline: %s" % ph.stop_timing(time_start))

    return 0
    def _run(self):

        ph.print_title("Two-step S2V-Registration and SRR Reconstruction")

        # Use linear spacing for alphas excluding the last alpha reserved
        # for the final SRR step
        alphas = np.linspace(self._alpha_range[0], self._alpha_range[1],
                             self._cycles)
        alphas = alphas[0:self._cycles]

        thresholds = np.linspace(self._threshold_range[0],
                                 self._threshold_range[1], self._cycles)
        thresholds = thresholds[0:self._cycles]

        s2vreg = SliceToVolumeRegistration(
            stacks=self._stacks,
            reference=self._reference,
            registration_method=self._registration_method,
            verbose=self._verbose,
            threshold_measure=self._threshold_measure,
            interleave=self._interleave,
        )

        reference = self._reference

        for cycle in range(0, self._cycles):

            # Slice-to-volume registration step
            s2vreg.set_reference(reference)
            s2vreg.set_print_prefix("Cycle %d/%d: " %
                                    (cycle + 1, self._cycles))
            if self._outlier_rejection:
                s2vreg.set_threshold(thresholds[cycle])
            if self._use_robust_registration and cycle == 0:
                s2vreg.set_s2v_smoothing(self._s2v_smoothing)
            else:
                s2vreg.set_s2v_smoothing(None)
            s2vreg.run()

            self._computational_time_registration += \
                s2vreg.get_computational_time()

            # SRR step
            if cycle < self._cycles - 1:
                # ---------------- Perform Image Reconstruction ---------------
                ph.print_subtitle("Volumetric Image Reconstruction")
                if isinstance(self._reconstruction_method,
                              sda.ScatteredDataApproximation):
                    self._reconstruction_method.set_sigma(alphas[cycle])
                else:
                    self._reconstruction_method.set_alpha(alphas[cycle])
                self._reconstruction_method.run()

                self._computational_time_reconstruction += \
                    self._reconstruction_method.get_computational_time()

                reference = self._reconstruction_method.get_reconstruction()

                # ------------------ Perform Image Mask SDA -------------------
                ph.print_subtitle("Volumetric Image Mask Reconstruction")
                SDA = sda.ScatteredDataApproximation(
                    self._stacks,
                    reference,
                    sigma=self._sigma_sda_mask,
                    sda_mask=True,
                )
                SDA.run()

                # reference contains updated mask based on SDA
                reference = SDA.get_reconstruction()

                # -------------------- Store Reconstruction -------------------
                filename = "Iter%d_%s" % (cycle + 1,
                                          self._reconstruction_method.
                                          get_setting_specific_filename())
                self._reconstructions.insert(
                    0, st.Stack.from_stack(reference, filename=filename))

                if self._verbose:
                    sitkh.show_stacks(self._reconstructions,
                                      segmentation=self._reference,
                                      viewer=self._viewer)
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    # Read input
    input_parser = InputArgparser(
        description="Script to study reconstruction parameters and their "
        "impact on the volumetric reconstruction quality. "
        "This script can only be used to sweep through one single parameter, "
        "e.g. the regularization parameter 'alpha'. ")
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks()
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_dir_input_mc()
    input_parser.add_dir_output(required=True)
    input_parser.add_reconstruction_space()
    input_parser.add_reference(
        help="Path to reference NIfTI image file. If given the volumetric "
        "reconstructed is performed in this physical space. "
        "Either a reconstruction space or a reference must be provided",
        required=False)
    input_parser.add_reference_mask(default=None)
    input_parser.add_study_name()
    input_parser.add_reconstruction_type(default="TK1L2")
    input_parser.add_measures(
        default=["PSNR", "MAE", "RMSE", "SSIM", "NCC", "NMI"])
    input_parser.add_tv_solver(default="PD")
    input_parser.add_iterations(default=50)
    input_parser.add_rho(default=0.1)
    input_parser.add_iter_max(default=10)
    input_parser.add_minimizer(default="lsmr")
    input_parser.add_log_config(default=1)
    input_parser.add_use_masks_srr(default=0)
    input_parser.add_verbose(default=1)
    input_parser.add_slice_thicknesses(default=None)
    input_parser.add_argument(
        "--append",
        "-append",
        action='store_true',
        help="If given, results are appended to previously executed parameter "
        "study with identical parameters and study name store in the output "
        "directory.")

    # Range for parameter sweeps
    input_parser.add_alphas(default=list(np.linspace(0.01, 0.5, 5)))
    input_parser.add_data_losses(default=["linear"]
                                 # default=["linear", "arctan"]
                                 )
    input_parser.add_data_loss_scales(default=[1]
                                      # default=[0.1, 0.5, 1.5]
                                      )

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.reference is None and args.reconstruction_space is None:
        raise IOError("Either reference (--reference) or reconstruction space "
                      "(--reconstruction-space) must be provided.")

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    data_reader = dr.MultipleImagesReader(
        file_paths=args.filenames,
        file_paths_masks=args.filenames_masks,
        suffix_mask=args.suffix_mask,
        dir_motion_correction=args.dir_input_mc,
        stacks_slice_thicknesses=args.slice_thicknesses,
    )

    data_reader.read_data()
    stacks = data_reader.get_data()
    ph.print_info("%d input stacks read for further processing" % len(stacks))

    if args.reference is not None:
        reference = st.Stack.from_filename(file_path=args.reference,
                                           file_path_mask=args.reference_mask,
                                           extract_slices=False)

        reconstruction_space = stacks[0].get_resampled_stack(reference.sitk)
        reconstruction_space = \
            reconstruction_space.get_stack_multiplied_with_mask()
        x_ref = sitk.GetArrayFromImage(reference.sitk).flatten()
        x_ref_mask = sitk.GetArrayFromImage(reference.sitk_mask).flatten()

    else:
        reconstruction_space = st.Stack.from_filename(
            file_path=args.reconstruction_space, extract_slices=False)
        reconstruction_space = stacks[0].get_resampled_stack(
            reconstruction_space.sitk)
        reconstruction_space = \
            reconstruction_space.get_stack_multiplied_with_mask()
        x_ref = None
        x_ref_mask = None

    # ----------------------------Set Up Parameters----------------------------
    parameters = {}
    parameters["alpha"] = args.alphas
    if len(args.data_losses) > 1:
        parameters["data_loss"] = args.data_losses
    if len(args.data_loss_scales) > 1:
        parameters["data_loss_scale"] = args.data_loss_scales

    # --------------------------Set Up Parameter Study-------------------------
    ph.print_title("Run Parameter Study")
    if args.study_name is None:
        name = args.reconstruction_type
    else:
        name = args.study_name

    reconstruction_info = {
        "shape": reconstruction_space.sitk.GetSize()[::-1],
        "origin": reconstruction_space.sitk.GetOrigin(),
        "spacing": reconstruction_space.sitk.GetSpacing(),
        "direction": reconstruction_space.sitk.GetDirection(),
    }

    # Create Tikhonov solver from which all information can be extracted
    # (also for other reconstruction types)
    tmp = tk.TikhonovSolver(
        stacks=stacks,
        reconstruction=reconstruction_space,
        alpha=args.alphas[0],
        iter_max=args.iter_max,
        data_loss=args.data_losses[0],
        data_loss_scale=args.data_loss_scales[0],
        reg_type="TK1",
        minimizer=args.minimizer,
        verbose=args.verbose,
        use_masks=args.use_masks_srr,
    )
    solver = tmp.get_solver()

    parameter_study_interface = \
        deconv_interface.DeconvolutionParameterStudyInterface(
            A=solver.get_A(),
            A_adj=solver.get_A_adj(),
            D=solver.get_B(),
            D_adj=solver.get_B_adj(),
            b=solver.get_b(),
            x0=solver.get_x0(),
            alpha=solver.get_alpha(),
            x_scale=solver.get_x_scale(),
            data_loss=solver.get_data_loss(),
            data_loss_scale=solver.get_data_loss_scale(),
            iter_max=solver.get_iter_max(),
            minimizer=solver.get_minimizer(),
            iterations=args.iterations,
            measures=args.measures,
            dimension=3,
            L2=16. / reconstruction_space.sitk.GetSpacing()[0]**2,
            reconstruction_type=args.reconstruction_type,
            rho=args.rho,
            dir_output=args.dir_output,
            parameters=parameters,
            name=name,
            reconstruction_info=reconstruction_info,
            x_ref=x_ref,
            x_ref_mask=x_ref_mask,
            tv_solver=args.tv_solver,
            verbose=args.verbose,
            append=args.append,
        )
    parameter_study_interface.set_up_parameter_study()
    parameter_study = parameter_study_interface.get_parameter_study()

    # Run parameter study
    parameter_study.run()

    print("\nComputational time for Deconvolution Parameter Study %s: %s" %
          (name, parameter_study.get_computational_time()))

    return 0
Beispiel #25
0
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Propagate image mask using rigid registration.", )
    input_parser.add_moving(required=True)
    input_parser.add_moving_mask(required=True)
    input_parser.add_fixed(required=True)
    input_parser.add_output(required=True)
    input_parser.add_v2v_method(
        option_string="--method",
        help="Registration method used for the registration (%s)." %
        (", or ".join(V2V_METHOD_OPTIONS)),
        default="RegAladin",
    )
    input_parser.add_option(
        option_string="--use-moving-mask",
        type=int,
        help="Turn on/off use of moving mask to constrain the registration.",
        default=0,
    )
    input_parser.add_dilation_radius(default=1)
    input_parser.add_verbose(default=0)
    input_parser.add_log_config(default=0)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if np.alltrue([not args.output.endswith(t) for t in ALLOWED_EXTENSIONS]):
        raise ValueError(
            "output filename invalid; allowed extensions are: %s" %
            ", ".join(ALLOWED_EXTENSIONS))

    if args.method not in V2V_METHOD_OPTIONS:
        raise ValueError("method must be in {%s}" %
                         (", ".join(V2V_METHOD_OPTIONS)))

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    stack = st.Stack.from_filename(
        file_path=args.fixed,
        extract_slices=False,
    )
    template = st.Stack.from_filename(
        file_path=args.moving,
        file_path_mask=args.moving_mask,
        extract_slices=False,
    )

    if args.method == "FLIRT":
        # Define search angle ranges for FLIRT in all three dimensions
        # search_angles = ["-searchr%s -%d %d" %
        #                  (x, args.search_angle, args.search_angle)
        #                  for x in ["x", "y", "z"]]
        # options = (" ").join(search_angles)
        # options += " -noresample"

        registration = regflirt.FLIRT(
            registration_type="Rigid",
            fixed=stack,
            moving=template,
            use_fixed_mask=False,
            use_moving_mask=args.use_moving_mask,
            # options=options,
            use_verbose=False,
        )
    else:
        registration = niftyreg.RegAladin(
            registration_type="Rigid",
            fixed=stack,
            moving=template,
            use_fixed_mask=False,
            use_moving_mask=args.use_moving_mask,
            # options="-ln 2",
            use_verbose=False,
        )

    try:
        registration.run()
    except RuntimeError as e:
        raise RuntimeError(
            "%s\n\n"
            "Have you tried running the script with '--use-moving-mask 0'?" %
            e)

    transform_sitk = registration.get_registration_transform_sitk()
    stack.sitk_mask = sitk.Resample(template.sitk_mask, stack.sitk_mask,
                                    transform_sitk, sitk.sitkNearestNeighbor,
                                    0, template.sitk_mask.GetPixelIDValue())
    if args.dilation_radius > 0:
        stack_mask_morpher = stmorph.StackMaskMorphologicalOperations.from_sitk_mask(
            mask_sitk=stack.sitk_mask,
            dilation_radius=args.dilation_radius,
            dilation_kernel="Ball",
            use_dilation_in_plane_only=True,
        )
        stack_mask_morpher.run_dilation()
        stack.sitk_mask = stack_mask_morpher.get_processed_mask_sitk()

    dw.DataWriter.write_mask(stack.sitk_mask, args.output)

    elapsed_time = ph.stop_timing(time_start)

    if args.verbose:
        ph.show_nifti(args.fixed, segmentation=args.output)

    ph.print_title("Summary")
    exe_file_info = os.path.basename(os.path.abspath(__file__)).split(".")[0]
    print("%s | Computational Time for Segmentation Propagation: %s" %
          (exe_file_info, elapsed_time))

    return 0
Beispiel #26
0
def main():

    time_start = ph.start_timing()

    flag_individual_cases_only = 1

    flag_batch_script = 0
    batch_ctr = [32]

    flag_correct_bias_field = 0
    # flag_correct_intensities = 0

    flag_collect_segmentations = 0
    flag_select_images_segmentations = 0

    flag_reconstruct_volume_subject_space = 0
    flag_reconstruct_volume_subject_space_irtk = 0
    flag_reconstruct_volume_subject_space_show_comparison = 0
    flag_register_to_template = 0
    flag_register_to_template_irtk = 0
    flag_show_srr_template_space = 0
    flag_reconstruct_volume_template_space = 0
    flag_collect_volumetric_reconstruction_results = 0
    flag_show_volumetric_reconstruction_results = 0

    flag_rsync_stuff = 0

    # Analysis
    flag_remove_failed_cases_for_analysis = 1
    flag_postop = 2  # 0... preop, 1...postop, 2... pre+postop

    flag_evaluate_image_similarities = 0
    flag_analyse_image_similarities = 1

    flag_evaluate_slice_residual_similarities = 0
    flag_analyse_slice_residual_similarities = 0

    flag_analyse_stacks = 0
    flag_analyse_qualitative_assessment = 0

    flag_collect_data_blinded_analysis = 0
    flag_anonymize_data_blinded_analysis = 0

    provide_comparison = 0
    intensity_correction = 1
    isotropic_resolution = 0.75
    alpha = 0.02
    outlier_rejection = 1
    threshold = 0.7
    threshold_first = 0.6

    # metric = "ANTSNeighborhoodCorrelation"
    # metric_radius = 5
    # multiresolution = 0

    prefix_srr = "srr_"
    prefix_srr_qa = "masked_"

    # ----------------------------------Set Up---------------------------------
    if flag_correct_bias_field:
        dir_batch = os.path.join(utils.DIR_BATCH_ROOT, "BiasFieldCorrection")
    elif flag_reconstruct_volume_subject_space:
        dir_batch = os.path.join(utils.DIR_BATCH_ROOT,
                                 "VolumetricReconstructionSubjectSpace")
    elif flag_register_to_template:
        dir_batch = os.path.join(utils.DIR_BATCH_ROOT,
                                 "VolumetricReconstructionRegisterToTemplate")
    elif flag_reconstruct_volume_template_space:
        dir_batch = os.path.join(utils.DIR_BATCH_ROOT,
                                 "VolumetricReconstructionTemplateSpace")
    else:
        dir_batch = os.path.join(utils.DIR_BATCH_ROOT, "foo")
    file_prefix_batch = os.path.join(dir_batch, "command")

    if flag_batch_script:
        verbose = 0
    else:
        verbose = 1

    data_reader = dr.ExcelSheetDataReader(utils.EXCEL_FILE)
    data_reader.read_data()
    cases = data_reader.get_data()

    if flag_analyse_qualitative_assessment:
        data_reader = dr.ExcelSheetQualitativeAssessmentReader(utils.QA_FILE)
        data_reader.read_data()
        qualitative_assessment = data_reader.get_data()

        statistical_evaluation = se.StatisticalEvaluation(
            qualitative_assessment)
        statistical_evaluation.run_tests(ref="seg_manual")
        ph.exit()

    cases_similarities = []
    cases_stacks = []

    if flag_individual_cases_only:
        N_cases = len(INDIVIDUAL_CASE_IDS)
    else:
        N_cases = len(cases.keys())

    i_case = 0
    for case_id in sorted(cases.keys()):
        if flag_individual_cases_only and case_id not in INDIVIDUAL_CASE_IDS:
            continue
        if not flag_analyse_image_similarities and \
                not flag_analyse_slice_residual_similarities:
            i_case += 1
            ph.print_title("%d/%d: %s" % (i_case, N_cases, case_id))

        if flag_rsync_stuff:
            dir_output = utils.get_directory_case_recon_seg_mode(
                case_id=case_id, recon_space="template_space", seg_mode="")

            dir_input = re.sub("Volumes/spina/",
                               "Volumes/medic-volumetric_res/SpinaBifida/",
                               dir_output)
            cmd = "rsync -avuhn --exclude 'motion_correction' %sseg_manual %s" % (
                dir_input, dir_output)
            ph.print_execution(cmd)
            # ph.execute_command(cmd)

        # -------------------------Correct Bias Field--------------------------
        if flag_correct_bias_field:
            filenames = utils.get_filenames_preprocessing_bias_field(case_id)
            paths_to_filenames = [
                os.path.join(utils.get_directory_case_original(case_id), f)
                for f in filenames
            ]
            dir_output = utils.get_directory_case_preprocessing(
                case_id, stage="01_N4ITK")

            # no image found matching the pattern
            if len(paths_to_filenames) == 0:
                continue

            cmd_args = []
            cmd_args.append("--filenames %s" % " ".join(paths_to_filenames))
            cmd_args.append("--dir-output %s" % dir_output)
            cmd_args.append("--prefix-output ''")
            cmd = "niftymic_correct_bias_field %s" % (" ").join(cmd_args)

            ph.execute_command(cmd,
                               flag_print_to_file=flag_batch_script,
                               path_to_file="%s%d.txt" %
                               (file_prefix_batch, ph.add_one(batch_ctr)))

        # # Skip case in case segmentations have not been provided yet
        # if not ph.directory_exists(utils.get_directory_case_segmentation(
        #         case_id, utils.SEGMENTATION_INIT, SEG_MODES[0])):
        #     continue

        # ------------------------Collect Segmentations------------------------
        if flag_collect_segmentations:
            # Skip case in case segmentations have been collected already
            if ph.directory_exists(
                    utils.get_directory_case_segmentation(
                        case_id, utils.SEGMENTATION_SELECTED, SEG_MODES[0])):
                ph.print_info("skipped")
                continue

            filenames = utils.get_segmented_image_filenames(
                case_id, subfolder=utils.SEGMENTATION_INIT)

            for i_seg_mode, seg_mode in enumerate(SEG_MODES):
                directory_selected = utils.get_directory_case_segmentation(
                    case_id, utils.SEGMENTATION_SELECTED, seg_mode)
                ph.create_directory(directory_selected)
                paths_to_filenames_init = [
                    os.path.join(
                        utils.get_directory_case_segmentation(
                            case_id, utils.SEGMENTATION_INIT, seg_mode), f)
                    for f in filenames
                ]
                paths_to_filenames_selected = [
                    os.path.join(directory_selected, f) for f in filenames
                ]
                for i in range(len(filenames)):
                    cmd = "cp -p %s %s" % (paths_to_filenames_init[i],
                                           paths_to_filenames_selected[i])
                    # ph.print_execution(cmd)
                    ph.execute_command(cmd)

        if flag_select_images_segmentations:
            filenames = utils.get_segmented_image_filenames(
                case_id, subfolder=utils.SEGMENTATION_SELECTED)
            paths_to_filenames = [
                os.path.join(
                    utils.get_directory_case_preprocessing(case_id,
                                                           stage="01_N4ITK"),
                    f) for f in filenames
            ]
            paths_to_filenames_masks = [
                os.path.join(
                    utils.get_directory_case_segmentation(
                        case_id, utils.SEGMENTATION_SELECTED, "seg_manual"), f)
                for f in filenames
            ]
            for i in range(len(filenames)):
                ph.show_niftis(
                    [paths_to_filenames[i]],
                    segmentation=paths_to_filenames_masks[i],
                    # viewer="fsleyes",
                )
                ph.pause()
                ph.killall_itksnap()

        # # -------------------------Correct Intensities-----------------------
        # if flag_correct_intensities:
        #     filenames = utils.get_segmented_image_filenames(case_id)
        #     paths_to_filenames_bias = [os.path.join(
        #         utils.get_directory_case_preprocessing(
        #             case_id, stage="01_N4ITK"), f) for f in filenames]
        #     print paths_to_filenames_bias

        # -----------------Reconstruct Volume in Subject Space-----------------
        if flag_reconstruct_volume_subject_space:

            filenames = utils.get_segmented_image_filenames(
                case_id, subfolder=utils.SEGMENTATION_SELECTED)
            # filenames = filenames[0:2]

            paths_to_filenames = [
                os.path.join(
                    utils.get_directory_case_preprocessing(case_id,
                                                           stage="01_N4ITK"),
                    f) for f in filenames
            ]

            # Estimate target stack
            target_stack_index = utils.get_target_stack_index(
                case_id, utils.SEGMENTATION_SELECTED, "seg_auto", filenames)

            for i, seg_mode in enumerate(SEG_MODES):
                # Get mask filenames
                paths_to_filenames_masks = [
                    os.path.join(
                        utils.get_directory_case_segmentation(
                            case_id, utils.SEGMENTATION_SELECTED, seg_mode), f)
                    for f in filenames
                ]

                if flag_reconstruct_volume_subject_space_irtk:
                    if seg_mode != "seg_manual":
                        continue
                    utils.export_irtk_call_to_workstation(
                        case_id=case_id,
                        filenames=filenames,
                        seg_mode=seg_mode,
                        isotropic_resolution=isotropic_resolution,
                        target_stack_index=target_stack_index,
                        kernel_mask_dilation=(15, 15, 4))

                else:
                    dir_output = utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="subject_space",
                        seg_mode=seg_mode)
                    # dir_output = "/tmp/foo"

                    cmd_args = []
                    cmd_args.append("--filenames %s" %
                                    " ".join(paths_to_filenames))
                    cmd_args.append("--filenames-masks %s" %
                                    " ".join(paths_to_filenames_masks))
                    cmd_args.append("--dir-output %s" % dir_output)
                    cmd_args.append("--use-masks-srr 0")
                    cmd_args.append("--isotropic-resolution %f" %
                                    isotropic_resolution)
                    cmd_args.append("--target-stack-index %d" %
                                    target_stack_index)
                    cmd_args.append("--intensity-correction %d" %
                                    intensity_correction)
                    cmd_args.append("--outlier-rejection %d" %
                                    outlier_rejection)
                    cmd_args.append("--threshold-first %f" % threshold_first)
                    cmd_args.append("--threshold %f" % threshold)
                    # cmd_args.append("--metric %s" % metric)
                    # cmd_args.append("--multiresolution %d" % multiresolution)
                    # cmd_args.append("--metric-radius %s" % metric_radius)
                    # if i > 0:
                    #     cmd_args.append("--reconstruction-space %s" % (
                    #         utils.get_path_to_recon(
                    #             utils.get_directory_case_recon_seg_mode(
                    #                 case_id, "seg_manual"))))
                    # cmd_args.append("--two-step-cycles 0")
                    cmd_args.append("--verbose %d" % verbose)
                    cmd_args.append("--provide-comparison %d" %
                                    provide_comparison)
                    # cmd_args.append("--iter-max 1")

                    cmd = "niftymic_reconstruct_volume %s" % (
                        " ").join(cmd_args)

                    ph.execute_command(
                        cmd,
                        flag_print_to_file=flag_batch_script,
                        path_to_file="%s%d.txt" %
                        (file_prefix_batch, ph.add_one(batch_ctr)))

        if flag_reconstruct_volume_subject_space_show_comparison:
            recon_paths = []
            for seg_mode in SEG_MODES:
                path_to_recon = utils.get_path_to_recon(
                    utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="subject_space",
                        seg_mode=seg_mode))
                recon_paths.append(path_to_recon)
            recon_path_irtk = os.path.join(
                utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="subject_space",
                    seg_mode="IRTK"), "IRTK_SRR.nii.gz")
            show_modes = list(SEG_MODES)
            if ph.file_exists(recon_path_irtk):
                recon_paths.append(recon_path_irtk)
                show_modes.append("irtk")
            ph.show_niftis(recon_paths)
            ph.print_info("Sequence: %s" % (" -- ").join(show_modes))
            ph.pause()
            ph.killall_itksnap()

        # -------------------------Register to template------------------------
        if flag_register_to_template:
            for seg_mode in SEG_MODES:

                cmd_args = []
                # register seg_auto-recon to template space
                if seg_mode == "seg_auto":

                    path_to_recon = utils.get_path_to_recon(
                        utils.get_directory_case_recon_seg_mode(
                            case_id=case_id,
                            recon_space="subject_space",
                            seg_mode=seg_mode))

                    template_stack_estimator = \
                        tse.TemplateStackEstimator.from_mask(
                            ph.append_to_filename(path_to_recon, "_mask"))
                    path_to_reference = \
                        template_stack_estimator.get_path_to_template()

                    dir_input_motion_correction = os.path.join(
                        utils.get_directory_case_recon_seg_mode(
                            case_id=case_id,
                            recon_space="subject_space",
                            seg_mode=seg_mode), "motion_correction")

                    dir_output = utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="template_space",
                        seg_mode=seg_mode)
                    # dir_output = "/home/mebner/tmp"
                    # # ------- DELETE -----
                    # dir_output = re.sub("data", "foo+1", dir_output)
                    # dir_output = re.sub(
                    #     "volumetric_reconstruction/20180126/template_space/seg_auto",
                    #     "", dir_output)
                    # # -------
                    # cmd_args.append("--use-fixed-mask 1")
                    cmd_args.append("--use-moving-mask 1")

                    # HACK
                    path_to_initial_transform = os.path.join(
                        utils.DIR_INPUT_ROOT_DATA, case_id,
                        "volumetric_reconstruction", "20180126",
                        "template_space", "seg_manual",
                        "registration_transform_sitk.txt")
                    cmd_args.append("--initial-transform %s" %
                                    path_to_initial_transform)
                    cmd_args.append("--use-flirt 0")
                    cmd_args.append("--use-regaladin 1")
                    cmd_args.append("--test-ap-flip 0")

                # register remaining recons to registered seg_auto-recon
                else:
                    path_to_reference = utils.get_path_to_recon(
                        utils.get_directory_case_recon_seg_mode(
                            case_id=case_id,
                            recon_space="template_space",
                            seg_mode="seg_auto"),
                        suffix="ResamplingToTemplateSpace",
                    )
                    path_to_initial_transform = os.path.join(
                        utils.get_directory_case_recon_seg_mode(
                            case_id=case_id,
                            recon_space="template_space",
                            seg_mode="seg_auto"),
                        "registration_transform_sitk.txt")

                    path_to_recon = utils.get_path_to_recon(
                        utils.get_directory_case_recon_seg_mode(
                            case_id=case_id,
                            recon_space="subject_space",
                            seg_mode=seg_mode))
                    dir_input_motion_correction = os.path.join(
                        utils.get_directory_case_recon_seg_mode(
                            case_id=case_id,
                            recon_space="subject_space",
                            seg_mode=seg_mode), "motion_correction")
                    dir_output = utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="template_space",
                        seg_mode=seg_mode)

                    cmd_args.append("--use-fixed-mask 0")
                    cmd_args.append("--use-moving-mask 0")
                    cmd_args.append("--initial-transform %s" %
                                    path_to_initial_transform)
                    cmd_args.append("--use-flirt 0")
                    cmd_args.append("--use-regaladin 1")
                    cmd_args.append("--test-ap-flip 0")

                cmd_args.append("--moving %s" % path_to_recon)
                cmd_args.append("--fixed %s" % path_to_reference)
                cmd_args.append("--dir-input %s" % dir_input_motion_correction)
                cmd_args.append("--dir-output %s" % dir_output)
                cmd_args.append("--write-transform 1")
                cmd_args.append("--verbose %d" % verbose)
                cmd = "niftymic_register_image %s" % (" ").join(cmd_args)

                ph.execute_command(cmd,
                                   flag_print_to_file=flag_batch_script,
                                   path_to_file="%s%d.txt" %
                                   (file_prefix_batch, ph.add_one(batch_ctr)))

        if flag_register_to_template_irtk:
            dir_input = utils.get_directory_case_recon_seg_mode(
                case_id=case_id, recon_space="subject_space", seg_mode="IRTK")
            dir_output = utils.get_directory_case_recon_seg_mode(
                case_id=case_id, recon_space="template_space", seg_mode="IRTK")
            path_to_recon = os.path.join(dir_input, "IRTK_SRR.nii.gz")
            path_to_reference = utils.get_path_to_recon(
                utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="template_space",
                    seg_mode="seg_manual"),
                suffix="ResamplingToTemplateSpace",
            )
            path_to_initial_transform = os.path.join(
                utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="template_space",
                    seg_mode="seg_manual"), "registration_transform_sitk.txt")

            cmd_args = []
            cmd_args.append("--fixed %s" % path_to_reference)
            cmd_args.append("--moving %s" % path_to_recon)
            cmd_args.append("--initial-transform %s" %
                            path_to_initial_transform)
            cmd_args.append("--use-fixed-mask 0")
            cmd_args.append("--use-moving-mask 0")
            cmd_args.append("--use-flirt 0")
            cmd_args.append("--use-regaladin 1")
            cmd_args.append("--test-ap-flip 0")
            cmd_args.append("--dir-output %s" % dir_output)
            cmd_args.append("--verbose %d" % verbose)
            cmd = "niftymic_register_image %s" % (" ").join(cmd_args)
            ph.execute_command(cmd)

        if flag_show_srr_template_space:
            recon_paths = []
            show_modes = list(SEG_MODES)
            # show_modes.append("IRTK")
            for seg_mode in show_modes:
                dir_input = utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="template_space",
                    seg_mode=seg_mode)
                # # ------- DELETE -----
                # dir_input = re.sub("data", "foo+1", dir_input)
                # dir_input = re.sub(
                #     "volumetric_reconstruction/20180126/template_space/seg_auto",
                #     "", dir_input)
                # # -------
                path_to_recon_space = utils.get_path_to_recon(
                    dir_input,
                    suffix="ResamplingToTemplateSpace",
                )
                recon_paths.append(path_to_recon_space)
            ph.show_niftis(recon_paths)
            ph.print_info("Sequence: %s" % (" -- ").join(show_modes))
            ph.pause()
            ph.killall_itksnap()

        # -----------------Reconstruct Volume in Template Space----------------
        if flag_reconstruct_volume_template_space:
            for seg_mode in SEG_MODES:
                path_to_recon_space = utils.get_path_to_recon(
                    utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="template_space",
                        seg_mode=seg_mode),
                    suffix="ResamplingToTemplateSpace",
                )
                dir_input = os.path.join(
                    utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="template_space",
                        seg_mode=seg_mode), "motion_correction")
                dir_output = utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="template_space",
                    seg_mode=seg_mode)
                # dir_output = os.path.join("/tmp/spina/template_space/%s-%s" % (
                #     case_id, seg_mode))

                cmd_args = []
                cmd_args.append("--dir-input %s" % dir_input)
                cmd_args.append("--dir-output %s" % dir_output)
                cmd_args.append("--reconstruction-space %s" %
                                path_to_recon_space)
                cmd_args.append("--alpha %s" % alpha)
                cmd_args.append("--verbose %s" % verbose)
                cmd_args.append("--use-masks-srr 0")

                # cmd_args.append("--minimizer L-BFGS-B")
                # cmd_args.append("--alpha 0.006")
                # cmd_args.append("--reconstruction-type HuberL2")
                # cmd_args.append("--data-loss arctan")
                # cmd_args.append("--iterations 5")
                # cmd_args.append("--data-loss-scale 0.7")

                cmd = "niftymic_reconstruct_volume_from_slices %s" % \
                    (" ").join(cmd_args)
                ph.execute_command(cmd,
                                   flag_print_to_file=flag_batch_script,
                                   path_to_file="%s%d.txt" %
                                   (file_prefix_batch, ph.add_one(batch_ctr)))

        # ----------------Collect SRR results in Template Space----------------
        if flag_collect_volumetric_reconstruction_results:
            directory = utils.get_directory_case_recon_summary(case_id)
            ph.create_directory(directory)

            # clear potentially existing files
            cmd = "rm -f %s/*.nii.gz" % (directory)
            ph.execute_command(cmd)

            # Collect SRRs
            for seg_mode in SEG_MODES:
                path_to_recon_src = utils.get_path_to_recon(
                    utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="template_space",
                        seg_mode=seg_mode), )
                path_to_recon = os.path.join(
                    directory, "%s%s.nii.gz" % (prefix_srr, seg_mode))

                cmd = "cp -p %s %s" % (path_to_recon_src, path_to_recon)
                ph.execute_command(cmd)

            # Collect IRTK recon
            path_to_recon_src = os.path.join(
                utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="template_space",
                    seg_mode="IRTK"),
                "IRTK_SRR_LinearResamplingToTemplateSpace.nii.gz")

            path_to_recon = os.path.join(directory,
                                         "%s%s.nii.gz" % (prefix_srr, "irtk"))

            cmd = "cp -p %s %s" % (path_to_recon_src, path_to_recon)
            ph.execute_command(cmd)

            # Collect evaluation mask
            path_to_recon = utils.get_path_to_recon(
                utils.get_directory_case_recon_seg_mode(
                    case_id=case_id,
                    recon_space="subject_space",
                    seg_mode="seg_auto"))

            template_stack_estimator = \
                tse.TemplateStackEstimator.from_mask(
                    ph.append_to_filename(path_to_recon, "_mask"))
            path_to_template = \
                template_stack_estimator.get_path_to_template()
            path_to_template_mask_src = ph.append_to_filename(
                path_to_template, "_mask_dil")
            path_to_template_mask = "%s/" % directory

            cmd = "cp -p %s %s" % (path_to_template_mask_src,
                                   path_to_template_mask)
            ph.execute_command(cmd)

        if flag_show_volumetric_reconstruction_results:
            dir_output = utils.get_directory_case_recon_summary(case_id)
            paths_to_recons = []
            for seg_mode in RECON_MODES:
                path_to_recon = os.path.join(
                    dir_output, "%s%s.nii.gz" % (prefix_srr, seg_mode))
                paths_to_recons.append(path_to_recon)
            path_to_mask = "%s/STA*.nii.gz" % dir_output
            cmd = ph.show_niftis(paths_to_recons, segmentation=path_to_mask)
            sitkh.write_executable_file([cmd], dir_output=dir_output)
            ph.pause()
            ph.killall_itksnap()

        # ---------------------Evaluate Image Similarities---------------------
        if flag_evaluate_image_similarities:
            dir_input = utils.get_directory_case_recon_summary(case_id)
            dir_output = utils.get_directory_case_recon_similarities(case_id)
            paths_to_recons = []
            for seg_mode in ["seg_auto", "detect", "irtk"]:
                path_to_recon = os.path.join(
                    dir_input, "%s%s.nii.gz" % (prefix_srr, seg_mode))
                paths_to_recons.append(path_to_recon)
            path_to_reference = os.path.join(
                dir_input, "%s%s.nii.gz" % (prefix_srr, "seg_manual"))
            path_to_reference_mask = utils.get_path_to_mask(dir_input)

            cmd_args = []
            cmd_args.append("--filenames %s" % " ".join(paths_to_recons))
            cmd_args.append("--reference %s" % path_to_reference)
            cmd_args.append("--reference-mask %s" % path_to_reference_mask)
            # cmd_args.append("--verbose 1")
            cmd_args.append("--dir-output %s" % dir_output)

            exe = re.sub("pyc", "py",
                         os.path.abspath(evaluate_image_similarity.__file__))
            cmd_args.insert(0, exe)

            # clear potentially existing files
            cmd = "rm -f %s/*.txt" % (dir_output)
            ph.execute_command(cmd)

            cmd = "python %s" % " ".join(cmd_args)
            ph.execute_command(cmd)

        # -----------------Evaluate Slice Residual Similarities----------------
        if flag_evaluate_slice_residual_similarities:

            path_to_reference_mask = utils.get_path_to_mask(
                utils.get_directory_case_recon_summary(case_id))

            dir_output_root = \
                utils.get_directory_case_slice_residual_similarities(case_id)

            # clear potentially existing files
            # cmd = "rm -f %s/*.txt" % (dir_output_root)
            # ph.execute_command(cmd)

            for seg_mode in SEG_MODES:
                dir_input = os.path.join(
                    utils.get_directory_case_recon_seg_mode(
                        case_id=case_id,
                        recon_space="template_space",
                        seg_mode=seg_mode,
                    ), "motion_correction")
                path_to_reference = os.path.join(
                    utils.get_directory_case_recon_summary(case_id),
                    "%s%s.nii.gz" % (prefix_srr, seg_mode))
                dir_output = os.path.join(dir_output_root, seg_mode)

                cmd_args = []
                cmd_args.append("--dir-input %s" % dir_input)
                cmd_args.append("--reference %s" % path_to_reference)
                cmd_args.append("--reference-mask %s" % path_to_reference_mask)
                cmd_args.append("--use-reference-mask 1")
                cmd_args.append("--use-slice-masks 0")
                # cmd_args.append("--verbose 1")
                cmd_args.append("--dir-output %s" % dir_output)

                exe = re.sub("pyc", "py", os.path.abspath(esrs.__file__))
                cmd_args.insert(0, exe)

                cmd = "python %s" % " ".join(cmd_args)
                ph.execute_command(cmd)

        # Collect data for blinded analysis
        if flag_collect_data_blinded_analysis:
            if flag_remove_failed_cases_for_analysis and case_id in RECON_FAILED_CASE_IDS:
                continue

            dir_input = utils.get_directory_case_recon_summary(case_id)
            # pattern = "STA([0-9]+)[_]mask.nii.gz"
            pattern = "STA([0-9]+)[_]mask_dil.nii.gz"
            p = re.compile(pattern)
            gw = [
                p.match(f).group(1) for f in os.listdir(dir_input)
                if p.match(f)
            ][0]

            dir_output = os.path.join(
                utils.get_directory_blinded_analysis(case_id, "open"), case_id)

            exe = re.sub("pyc", "py", os.path.abspath(mswm.__file__))

            recons = []

            for seg_mode in RECON_MODES:
                path_to_recon = os.path.join(
                    dir_input, "%s%s.nii.gz" % (prefix_srr, seg_mode))

                cmd_args = []
                cmd_args.append("--filename %s" % path_to_recon)
                cmd_args.append("--gestational-age %s" % gw)
                cmd_args.append("--dir-output %s" % dir_output)
                cmd_args.append("--prefix-output %s" % prefix_srr_qa)
                cmd_args.append("--verbose 0")
                cmd_args.insert(0, exe)

                cmd = "python %s" % " ".join(cmd_args)
                # ph.execute_command(cmd)

                recon = "%s%s" % (prefix_srr_qa,
                                  os.path.basename(path_to_recon))
                recons.append(recon)
            ph.write_show_niftis_exe(recons, dir_output)

        if flag_anonymize_data_blinded_analysis:
            dir_input = os.path.join(
                utils.get_directory_blinded_analysis(case_id, "open"), case_id)
            dir_output_dictionaries = utils.get_directory_anonymized_dictionares(
                case_id)
            dir_output_anonymized_images = os.path.join(
                utils.get_directory_blinded_analysis(case_id, "anonymized"),
                case_id)

            if not ph.directory_exists(dir_input):
                continue
            ph.create_directory(dir_output_dictionaries)
            ph.create_directory(dir_output_anonymized_images)

            data_anonymizer = da.DataAnonymizer()
            # Create random dictionary (only required once)
            # data_anonymizer.set_prefix_identifiers("%s_" % case_id)
            # data_anonymizer.read_nifti_filenames_from_directory(dir_input)
            # data_anonymizer.generate_identifiers()
            # data_anonymizer.generate_randomized_dictionary()
            # data_anonymizer.write_dictionary(
            #     dir_output_dictionaries, "dictionary_%s" % case_id)

            # Read dictionary
            data_anonymizer.read_dictionary(dir_output_dictionaries,
                                            "dictionary_%s" % case_id)

            # Anonymize files
            if 0:
                ph.clear_directory(dir_output_anonymized_images)
                data_anonymizer.anonymize_files(dir_input,
                                                dir_output_anonymized_images)

                # Write executable script
                filenames = [
                    "%s.nii.gz" % f
                    for f in sorted(data_anonymizer.get_identifiers())
                ]
                ph.write_show_niftis_exe(filenames,
                                         dir_output_anonymized_images)

            # Reveal anonymized files
            if 1:
                filenames = data_anonymizer.reveal_anonymized_files(
                    dir_output_anonymized_images)
                filenames = sorted(["%s" % f for f in filenames])
                ph.write_show_niftis_exe(filenames,
                                         dir_output_anonymized_images)

            # Reveal additional, original files
            # data_anonymizer.reveal_original_files(dir_output)

            # relative_directory = re.sub(
            #     utils.get_directory_blinded_analysis(case_id, "anonymized"),
            #     ".",
            #     dir_output_anonymized_images)
            # paths_to_filenames = [os.path.join(
            #     relative_directory, f) for f in filenames]

        # ---------------------Analyse Image Similarities---------------------
        if flag_analyse_image_similarities or \
                flag_analyse_slice_residual_similarities or \
                flag_analyse_stacks:
            if flag_remove_failed_cases_for_analysis:
                if case_id in RECON_FAILED_CASE_IDS:
                    continue
            if cases[case_id]['postrep'] == flag_postop or flag_postop == 2:
                cases_similarities.append(case_id)
                cases_stacks.append(
                    utils.get_segmented_image_filenames(
                        case_id,
                        # subfolder=utils.SEGMENTATION_INIT,
                        subfolder=utils.SEGMENTATION_SELECTED,
                    ))

        dir_output_analysis = os.path.join(
            # "/Users/mebner/UCL/UCL/Publications",
            "/home/mebner/Dropbox/UCL/Publications",
            "2018_MICCAI/brain_reconstruction_paper")

    if flag_analyse_image_similarities:
        dir_inputs = []
        filename = "image_similarities_postop%d.txt" % flag_postop
        for case_id in cases_similarities:
            dir_inputs.append(
                utils.get_directory_case_recon_similarities(case_id))
        cmd_args = []
        cmd_args.append("--dir-inputs %s" % " ".join(dir_inputs))
        cmd_args.append("--dir-output %s" % dir_output_analysis)
        cmd_args.append("--filename %s" % filename)

        exe = re.sub("pyc", "py",
                     os.path.abspath(src.analyse_image_similarities.__file__))
        cmd_args.insert(0, exe)

        cmd = "python %s" % " ".join(cmd_args)
        ph.execute_command(cmd)

    if flag_analyse_slice_residual_similarities:
        dir_inputs = []
        filename = "slice_residuals_postop%d.txt" % flag_postop
        for case_id in cases_similarities:
            dir_inputs.append(
                utils.get_directory_case_slice_residual_similarities(case_id))
        cmd_args = []
        cmd_args.append("--dir-inputs %s" % " ".join(dir_inputs))
        cmd_args.append("--subfolder %s" % " ".join(SEG_MODES))
        cmd_args.append("--dir-output %s" % dir_output_analysis)
        cmd_args.append("--filename %s" % filename)

        exe = re.sub(
            "pyc", "py",
            os.path.abspath(src.analyse_slice_residual_similarities.__file__))
        cmd_args.insert(0, exe)

        cmd = "python %s" % " ".join(cmd_args)
        # print len(cases_similarities)
        # print cases_similarities
        ph.execute_command(cmd)

    if flag_analyse_stacks:
        cases_stacks_N = [len(s) for s in cases_stacks]
        ph.print_subtitle("%d cases -- Number of stacks" % len(cases_stacks))
        ph.print_info("min: %g" % np.min(cases_stacks_N))
        ph.print_info("mean: %g" % np.mean(cases_stacks_N))
        ph.print_info("median: %g" % np.median(cases_stacks_N))
        ph.print_info("max: %g" % np.max(cases_stacks_N))

    elapsed_time = ph.stop_timing(time_start)
    ph.print_title("Summary")
    print("Computational Time for Pipeline: %s" % (elapsed_time))

    return 0
Beispiel #27
0
def main():

    input_parser = InputArgparser(description="Convert NIfTI to DICOM image", )
    input_parser.add_filename(required=True)
    input_parser.add_option(
        option_string="--template",
        type=str,
        required=True,
        help="Template DICOM to extract relevant DICOM tags.",
    )
    input_parser.add_dir_output(required=True)
    input_parser.add_label(
        help="Label used for series description of DICOM output.",
        default="SRR_NiftyMIC")
    input_parser.add_argument(
        "--volume",
        "-volume",
        action='store_true',
        help="If given, the output DICOM file is combined as 3D volume")
    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    # Prepare for final DICOM output
    ph.create_directory(args.dir_output)

    if args.volume:
        dir_output_2d_slices = os.path.join(DIR_TMP, "dicom_slices")
    else:
        dir_output_2d_slices = os.path.join(args.dir_output, args.label)
    ph.create_directory(dir_output_2d_slices, delete_files=True)

    # read NiftyMIC version (if available)
    data_reader = dr.ImageHeaderReader(args.filename)
    data_reader.read_data()
    niftymic_version = data_reader.get_niftymic_version()
    if niftymic_version is None:
        niftymic_version = "NiftyMIC"
    else:
        niftymic_version = "NiftyMIC-v%s" % niftymic_version

    # Create set of 2D DICOM slices from 3D NIfTI image
    # (correct image orientation!)
    ph.print_title("Create set of 2D DICOM slices from 3D NIfTI image")
    cmd_args = ["nifti2dicom"]
    cmd_args.append("-i '%s'" % args.filename)
    cmd_args.append("-o '%s'" % dir_output_2d_slices)
    cmd_args.append("-d '%s'" % args.template)
    cmd_args.append("--prefix ''")
    cmd_args.append("--seriesdescription '%s'" % args.label)
    cmd_args.append("--accessionnumber '%s'" % ACCESSION_NUMBER)
    cmd_args.append("--seriesnumber '%s'" % SERIES_NUMBER)
    cmd_args.append("--institutionname '%s'" % IMAGE_COMMENTS)

    # Overwrite default "nifti2dicom" tags which would be added otherwise
    # (no deletion/update with empty '' sufficient to overwrite them)
    cmd_args.append("--manufacturersmodelname '%s'" % "NiftyMIC")
    cmd_args.append("--protocolname '%s'" % niftymic_version)

    cmd_args.append("-y")
    ph.execute_command(" ".join(cmd_args))

    if args.volume:
        path_to_output = os.path.join(args.dir_output, "%s.dcm" % args.label)
        # Combine set of 2D DICOM slices to form 3D DICOM image
        # (image orientation stays correct)
        ph.print_title("Combine set of 2D DICOM slices to form 3D DICOM image")
        cmd_args = ["medcon"]
        cmd_args.append("-f '%s'/*.dcm" % dir_output_2d_slices)
        cmd_args.append("-o '%s'" % path_to_output)
        cmd_args.append("-c dicom")
        cmd_args.append("-stack3d")
        cmd_args.append("-n")
        cmd_args.append("-qc")
        cmd_args.append("-w")
        ph.execute_command(" ".join(cmd_args))

        # Update all relevant DICOM tags accordingly
        ph.print_title("Update all relevant DICOM tags accordingly")
        print("")
        dataset_template = pydicom.dcmread(args.template)
        dataset = pydicom.dcmread(path_to_output)

        # Copy tags from template (to guarantee grouping with original data)
        update_dicom_tags = {}
        for tag in COPY_DICOM_TAGS:
            try:
                update_dicom_tags[tag] = getattr(dataset_template, tag)
            except:
                update_dicom_tags[tag] = ""

        # Additional tags
        update_dicom_tags["SeriesDescription"] = args.label
        update_dicom_tags["InstitutionName"] = institution_name
        update_dicom_tags["ImageComments"] = IMAGE_COMMENTS
        update_dicom_tags["AccessionNumber"] = ACCESSION_NUMBER
        update_dicom_tags["SeriesNumber"] = SERIES_NUMBER

        for tag in sorted(update_dicom_tags.keys()):
            value = update_dicom_tags[tag]
            setattr(dataset, tag, value)
            ph.print_info("%s: '%s'" % (tag, value))

        dataset.save_as(path_to_output)
        print("")
        ph.print_info("3D DICOM image written to '%s'" % path_to_output)

    else:
        ph.print_info("DICOM images written to '%s'" % dir_output_2d_slices)

    return 0
def main():

    time_start = ph.start_timing()

    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Perform (linear) intensity correction across "
        "stacks/images given a reference stack/image", )
    input_parser.add_filenames(required=True)
    input_parser.add_dir_output(required=True)
    input_parser.add_reference(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_search_angle(default=180)
    input_parser.add_prefix_output(default="IC_")
    input_parser.add_log_config(default=1)
    input_parser.add_option(
        option_string="--registration",
        type=int,
        help="Turn on/off registration from image to reference prior to "
        "intensity correction.",
        default=0)
    input_parser.add_verbose(default=0)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    if args.reference in args.filenames:
        args.filenames.remove(args.reference)

    # Read data
    data_reader = dr.MultipleImagesReader(args.filenames,
                                          suffix_mask=args.suffix_mask,
                                          extract_slices=False)
    data_reader.read_data()
    stacks = data_reader.get_data()

    data_reader = dr.MultipleImagesReader([args.reference],
                                          suffix_mask=args.suffix_mask,
                                          extract_slices=False)
    data_reader.read_data()
    reference = data_reader.get_data()[0]

    if args.registration:
        # Define search angle ranges for FLIRT in all three dimensions
        search_angles = [
            "-searchr%s -%d %d" % (x, args.search_angle, args.search_angle)
            for x in ["x", "y", "z"]
        ]
        search_angles = (" ").join(search_angles)
        registration = regflirt.FLIRT(
            moving=reference,
            registration_type="Rigid",
            use_fixed_mask=True,
            use_moving_mask=True,
            options=search_angles,
            use_verbose=False,
        )

    # Perform Intensity Correction
    ph.print_title("Perform Intensity Correction")
    intensity_corrector = ic.IntensityCorrection(
        use_reference_mask=True,
        use_individual_slice_correction=False,
        prefix_corrected=args.prefix_output,
        use_verbose=False,
    )
    stacks_corrected = [None] * len(stacks)
    for i, stack in enumerate(stacks):
        if args.registration:
            ph.print_info("Image %d/%d: Registration ... " %
                          (i + 1, len(stacks)),
                          newline=False)
            registration.set_fixed(stack)
            registration.run()
            transform_sitk = registration.get_registration_transform_sitk()
            stack.update_motion_correction(transform_sitk)
            print("done")

        ph.print_info("Image %d/%d: Intensity Correction ... " %
                      (i + 1, len(stacks)),
                      newline=False)

        ref = reference.get_resampled_stack(stack.sitk)
        ref = st.Stack.from_sitk_image(image_sitk=ref.sitk,
                                       image_sitk_mask=stack.sitk_mask *
                                       ref.sitk_mask,
                                       filename=reference.get_filename())
        intensity_corrector.set_stack(stack)
        intensity_corrector.set_reference(ref)
        intensity_corrector.run_linear_intensity_correction()
        # intensity_corrector.run_affine_intensity_correction()
        stacks_corrected[i] = \
            intensity_corrector.get_intensity_corrected_stack()
        print("done (c1 = %g) " %
              intensity_corrector.get_intensity_correction_coefficients())

        # Write Data
        stacks_corrected[i].write(args.dir_output,
                                  write_mask=True,
                                  suffix_mask=args.suffix_mask)

        if args.verbose:
            sitkh.show_stacks(
                [
                    reference,
                    stacks_corrected[i],
                    # stacks[i],
                ],
                segmentation=stacks_corrected[i])
            # ph.pause()

    # Write reference too (although not intensity corrected)
    reference.write(args.dir_output,
                    filename=args.prefix_output + reference.get_filename(),
                    write_mask=True,
                    suffix_mask=args.suffix_mask)

    elapsed_time = ph.stop_timing(time_start)

    ph.print_title("Summary")
    print("Computational Time for Intensity Correction(s): %s" %
          (elapsed_time))

    return 0
Beispiel #29
0
def main():

    time_start = ph.start_timing()

    # Set print options for numpy
    np.set_printoptions(precision=3)

    input_parser = InputArgparser(
        description="Volumetric MRI reconstruction framework to reconstruct "
        "an isotropic, high-resolution 3D volume from multiple stacks of 2D "
        "slices with motion correction. The resolution of the computed "
        "Super-Resolution Reconstruction (SRR) is given by the in-plane "
        "spacing of the selected target stack. A region of interest can be "
        "specified by providing a mask for the selected target stack. Only "
        "this region will then be reconstructed by the SRR algorithm which "
        "can substantially reduce the computational time.",
    )
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks()
    input_parser.add_output(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_target_stack(default=None)
    input_parser.add_search_angle(default=45)
    input_parser.add_multiresolution(default=0)
    input_parser.add_shrink_factors(default=[3, 2, 1])
    input_parser.add_smoothing_sigmas(default=[1.5, 1, 0])
    input_parser.add_sigma(default=1)
    input_parser.add_reconstruction_type(default="TK1L2")
    input_parser.add_iterations(default=15)
    input_parser.add_alpha(default=0.015)
    input_parser.add_alpha_first(default=0.2)
    input_parser.add_iter_max(default=10)
    input_parser.add_iter_max_first(default=5)
    input_parser.add_dilation_radius(default=3)
    input_parser.add_extra_frame_target(default=10)
    input_parser.add_bias_field_correction(default=0)
    input_parser.add_intensity_correction(default=1)
    input_parser.add_isotropic_resolution(default=1)
    input_parser.add_log_config(default=1)
    input_parser.add_subfolder_motion_correction()
    input_parser.add_write_motion_correction(default=1)
    input_parser.add_verbose(default=0)
    input_parser.add_two_step_cycles(default=3)
    input_parser.add_use_masks_srr(default=0)
    input_parser.add_boundary_stacks(default=[10, 10, 0])
    input_parser.add_metric(default="Correlation")
    input_parser.add_metric_radius(default=10)
    input_parser.add_reference()
    input_parser.add_reference_mask()
    input_parser.add_outlier_rejection(default=1)
    input_parser.add_threshold_first(default=0.5)
    input_parser.add_threshold(default=0.8)
    input_parser.add_interleave(default=3)
    input_parser.add_slice_thicknesses(default=None)
    input_parser.add_viewer(default="itksnap")
    input_parser.add_v2v_method(default="RegAladin")
    input_parser.add_argument(
        "--v2v-robust", "-v2v-robust",
        action='store_true',
        help="If given, a more robust volume-to-volume registration step is "
        "performed, i.e. four rigid registrations are performed using four "
        "rigid transform initializations based on "
        "principal component alignment of associated masks."
    )
    input_parser.add_argument(
        "--s2v-hierarchical", "-s2v-hierarchical",
        action='store_true',
        help="If given, a hierarchical approach for the first slice-to-volume "
        "registration cycle is used, i.e. sub-packages defined by the "
        "specified interleave (--interleave) are registered until each "
        "slice is registered independently."
    )
    input_parser.add_argument(
        "--sda", "-sda",
        action='store_true',
        help="If given, the volumetric reconstructions are performed using "
        "Scattered Data Approximation (Vercauteren et al., 2006). "
        "'alpha' is considered the final 'sigma' for the "
        "iterative adjustment. "
        "Recommended value is, e.g., --alpha 0.8"
    )
    input_parser.add_option(
        option_string="--transforms-history",
        type=int,
        help="Write entire history of applied slice motion correction "
        "transformations to motion correction output directory",
        default=0,
    )

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    rejection_measure = "NCC"
    threshold_v2v = -2  # 0.3
    debug = False

    if args.v2v_method not in V2V_METHOD_OPTIONS:
        raise ValueError("v2v-method must be in {%s}" % (
            ", ".join(V2V_METHOD_OPTIONS)))

    if np.alltrue([not args.output.endswith(t) for t in ALLOWED_EXTENSIONS]):
        raise ValueError(
            "output filename invalid; allowed extensions are: %s" %
            ", ".join(ALLOWED_EXTENSIONS))

    if args.alpha_first < args.alpha and not args.sda:
        raise ValueError("It must hold alpha-first >= alpha")

    if args.threshold_first > args.threshold:
        raise ValueError("It must hold threshold-first <= threshold")

    dir_output = os.path.dirname(args.output)
    ph.create_directory(dir_output)

    if args.log_config:
        input_parser.log_config(os.path.abspath(__file__))

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")
    data_reader = dr.MultipleImagesReader(
        file_paths=args.filenames,
        file_paths_masks=args.filenames_masks,
        suffix_mask=args.suffix_mask,
        stacks_slice_thicknesses=args.slice_thicknesses,
    )

    if len(args.boundary_stacks) is not 3:
        raise IOError(
            "Provide exactly three values for '--boundary-stacks' to define "
            "cropping in i-, j-, and k-dimension of the input stacks")

    data_reader.read_data()
    stacks = data_reader.get_data()
    ph.print_info("%d input stacks read for further processing" % len(stacks))

    if all(s.is_unity_mask() is True for s in stacks):
        ph.print_warning("No mask is provided! "
                         "Generated reconstruction space may be very big!")
        ph.print_warning("Consider using a mask to speed up computations")

        # args.extra_frame_target = 0
        # ph.wrint_warning("Overwritten: extra-frame-target set to 0")

    # Specify target stack for intensity correction and reconstruction space
    if args.target_stack is None:
        target_stack_index = 0
    else:
        try:
            target_stack_index = args.filenames.index(args.target_stack)
        except ValueError as e:
            raise ValueError(
                "--target-stack must correspond to an image as provided by "
                "--filenames")

    # ---------------------------Data Preprocessing---------------------------
    ph.print_title("Data Preprocessing")

    segmentation_propagator = segprop.SegmentationPropagation(
        # registration_method=regflirt.FLIRT(use_verbose=args.verbose),
        # registration_method=niftyreg.RegAladin(use_verbose=False),
        dilation_radius=args.dilation_radius,
        dilation_kernel="Ball",
    )

    data_preprocessing = dp.DataPreprocessing(
        stacks=stacks,
        segmentation_propagator=segmentation_propagator,
        use_cropping_to_mask=True,
        use_N4BiasFieldCorrector=args.bias_field_correction,
        target_stack_index=target_stack_index,
        boundary_i=args.boundary_stacks[0],
        boundary_j=args.boundary_stacks[1],
        boundary_k=args.boundary_stacks[2],
        unit="mm",
    )
    data_preprocessing.run()
    time_data_preprocessing = data_preprocessing.get_computational_time()

    # Get preprocessed stacks
    stacks = data_preprocessing.get_preprocessed_stacks()

    # Define reference/target stack for registration and reconstruction
    if args.reference is not None:
        reference = st.Stack.from_filename(
            file_path=args.reference,
            file_path_mask=args.reference_mask,
            extract_slices=False)

    else:
        reference = st.Stack.from_stack(stacks[target_stack_index])

    # ------------------------Volume-to-Volume Registration--------------------
    if len(stacks) > 1:

        if args.v2v_method == "FLIRT":
            # Define search angle ranges for FLIRT in all three dimensions
            search_angles = ["-searchr%s -%d %d" %
                             (x, args.search_angle, args.search_angle)
                             for x in ["x", "y", "z"]]
            options = (" ").join(search_angles)
            # options += " -noresample"

            vol_registration = regflirt.FLIRT(
                registration_type="Rigid",
                use_fixed_mask=True,
                use_moving_mask=True,
                options=options,
                use_verbose=False,
            )
        else:
            vol_registration = niftyreg.RegAladin(
                registration_type="Rigid",
                use_fixed_mask=True,
                use_moving_mask=True,
                # options="-ln 2 -voff",
                use_verbose=False,
            )
        v2vreg = pipeline.VolumeToVolumeRegistration(
            stacks=stacks,
            reference=reference,
            registration_method=vol_registration,
            verbose=debug,
            robust=args.v2v_robust,
        )
        v2vreg.run()
        stacks = v2vreg.get_stacks()
        time_registration = v2vreg.get_computational_time()

    else:
        time_registration = ph.get_zero_time()

    # ---------------------------Intensity Correction--------------------------
    if args.intensity_correction:
        ph.print_title("Intensity Correction")
        intensity_corrector = ic.IntensityCorrection()
        intensity_corrector.use_individual_slice_correction(False)
        intensity_corrector.use_reference_mask(True)
        intensity_corrector.use_stack_mask(True)
        intensity_corrector.use_verbose(False)

        for i, stack in enumerate(stacks):
            if i == target_stack_index:
                ph.print_info("Stack %d (%s): Reference image. Skipped." % (
                    i + 1, stack.get_filename()))
                continue
            else:
                ph.print_info("Stack %d (%s): Intensity Correction ... " % (
                    i + 1, stack.get_filename()), newline=False)
            intensity_corrector.set_stack(stack)
            intensity_corrector.set_reference(
                stacks[target_stack_index].get_resampled_stack(
                    resampling_grid=stack.sitk,
                    interpolator="NearestNeighbor",
                ))
            intensity_corrector.run_linear_intensity_correction()
            stacks[i] = intensity_corrector.get_intensity_corrected_stack()
            print("done (c1 = %g) " %
                  intensity_corrector.get_intensity_correction_coefficients())

    # ---------------------------Create first volume---------------------------
    time_tmp = ph.start_timing()

    # Isotropic resampling to define HR target space
    ph.print_title("Reconstruction Space Generation")
    HR_volume = reference.get_isotropically_resampled_stack(
        resolution=args.isotropic_resolution)
    ph.print_info(
        "Isotropic reconstruction space with %g mm resolution is created" %
        HR_volume.sitk.GetSpacing()[0])

    if args.reference is None:
        # Create joint image mask in target space
        joint_image_mask_builder = imb.JointImageMaskBuilder(
            stacks=stacks,
            target=HR_volume,
            dilation_radius=1,
        )
        joint_image_mask_builder.run()
        HR_volume = joint_image_mask_builder.get_stack()
        ph.print_info(
            "Isotropic reconstruction space is centered around "
            "joint stack masks. ")

        # Crop to space defined by mask (plus extra margin)
        HR_volume = HR_volume.get_cropped_stack_based_on_mask(
            boundary_i=args.extra_frame_target,
            boundary_j=args.extra_frame_target,
            boundary_k=args.extra_frame_target,
            unit="mm",
        )

        # Create first volume
        # If outlier rejection is activated, eliminate obvious outliers early
        # from stack and re-run SDA to get initial volume without them
        ph.print_title("First Estimate of HR Volume")
        if args.outlier_rejection and threshold_v2v > -1:
            ph.print_subtitle("SDA Approximation")
            SDA = sda.ScatteredDataApproximation(
                stacks, HR_volume, sigma=args.sigma)
            SDA.run()
            HR_volume = SDA.get_reconstruction()

            # Identify and reject outliers
            ph.print_subtitle("Eliminate slice outliers (%s < %g)" % (
                rejection_measure, threshold_v2v))
            outlier_rejector = outre.OutlierRejector(
                stacks=stacks,
                reference=HR_volume,
                threshold=threshold_v2v,
                measure=rejection_measure,
                verbose=True,
            )
            outlier_rejector.run()
            stacks = outlier_rejector.get_stacks()

        ph.print_subtitle("SDA Approximation Image")
        SDA = sda.ScatteredDataApproximation(
            stacks, HR_volume, sigma=args.sigma)
        SDA.run()
        HR_volume = SDA.get_reconstruction()

        ph.print_subtitle("SDA Approximation Image Mask")
        SDA = sda.ScatteredDataApproximation(
            stacks, HR_volume, sigma=args.sigma, sda_mask=True)
        SDA.run()
        # HR volume contains updated mask based on SDA
        HR_volume = SDA.get_reconstruction()

        HR_volume.set_filename(SDA.get_setting_specific_filename())

    time_reconstruction = ph.stop_timing(time_tmp)

    if args.verbose:
        tmp = list(stacks)
        tmp.insert(0, HR_volume)
        sitkh.show_stacks(tmp, segmentation=HR_volume, viewer=args.viewer)

    # -----------Two-step Slice-to-Volume Registration-Reconstruction----------
    if args.two_step_cycles > 0:

        # Slice-to-volume registration set-up
        if args.metric == "ANTSNeighborhoodCorrelation":
            metric_params = {"radius": args.metric_radius}
        else:
            metric_params = None
        registration = regsitk.SimpleItkRegistration(
            moving=HR_volume,
            use_fixed_mask=True,
            use_moving_mask=True,
            interpolator="Linear",
            metric=args.metric,
            metric_params=metric_params,
            use_multiresolution_framework=args.multiresolution,
            shrink_factors=args.shrink_factors,
            smoothing_sigmas=args.smoothing_sigmas,
            initializer_type="SelfGEOMETRY",
            optimizer="ConjugateGradientLineSearch",
            optimizer_params={
                "learningRate": 1,
                "numberOfIterations": 100,
                "lineSearchUpperLimit": 2,
            },
            scales_estimator="Jacobian",
            use_verbose=debug,
        )

        # Volumetric reconstruction set-up
        if args.sda:
            recon_method = sda.ScatteredDataApproximation(
                stacks,
                HR_volume,
                sigma=args.sigma,
                use_masks=args.use_masks_srr,
            )
            alpha_range = [args.sigma, args.alpha]
        else:
            recon_method = tk.TikhonovSolver(
                stacks=stacks,
                reconstruction=HR_volume,
                reg_type="TK1",
                minimizer="lsmr",
                alpha=args.alpha_first,
                iter_max=np.min([args.iter_max_first, args.iter_max]),
                verbose=True,
                use_masks=args.use_masks_srr,
            )
            alpha_range = [args.alpha_first, args.alpha]

        # Define the regularization parameters for the individual
        # reconstruction steps in the two-step cycles
        alphas = np.linspace(
            alpha_range[0], alpha_range[1], args.two_step_cycles)

        # Define outlier rejection threshold after each S2V-reg step
        thresholds = np.linspace(
            args.threshold_first, args.threshold, args.two_step_cycles)

        two_step_s2v_reg_recon = \
            pipeline.TwoStepSliceToVolumeRegistrationReconstruction(
                stacks=stacks,
                reference=HR_volume,
                registration_method=registration,
                reconstruction_method=recon_method,
                cycles=args.two_step_cycles,
                alphas=alphas[0:args.two_step_cycles - 1],
                outlier_rejection=args.outlier_rejection,
                threshold_measure=rejection_measure,
                thresholds=thresholds,
                interleave=args.interleave,
                viewer=args.viewer,
                verbose=args.verbose,
                use_hierarchical_registration=args.s2v_hierarchical,
            )
        two_step_s2v_reg_recon.run()
        HR_volume_iterations = \
            two_step_s2v_reg_recon.get_iterative_reconstructions()
        time_registration += \
            two_step_s2v_reg_recon.get_computational_time_registration()
        time_reconstruction += \
            two_step_s2v_reg_recon.get_computational_time_reconstruction()
        stacks = two_step_s2v_reg_recon.get_stacks()

    # no two-step s2v-registration/reconstruction iterations
    else:
        HR_volume_iterations = []

    # Write motion-correction results
    ph.print_title("Write Motion Correction Results")
    if args.write_motion_correction:
        dir_output_mc = os.path.join(
            dir_output, args.subfolder_motion_correction)
        ph.clear_directory(dir_output_mc)

        for stack in stacks:
            stack.write(
                dir_output_mc,
                write_stack=False,
                write_mask=False,
                write_slices=False,
                write_transforms=True,
                write_transforms_history=args.transforms_history,
            )

        if args.outlier_rejection:
            deleted_slices_dic = {}
            for i, stack in enumerate(stacks):
                deleted_slices = stack.get_deleted_slice_numbers()
                deleted_slices_dic[stack.get_filename()] = deleted_slices

            # check whether any stack was removed entirely
            stacks0 = data_preprocessing.get_preprocessed_stacks()
            if len(stacks) != len(stacks0):
                stacks_remain = [s.get_filename() for s in stacks]
                for stack in stacks0:
                    if stack.get_filename() in stacks_remain:
                        continue

                    # add info that all slices of this stack were rejected
                    deleted_slices = [
                        slice.get_slice_number()
                        for slice in stack.get_slices()
                    ]
                    deleted_slices_dic[stack.get_filename()] = deleted_slices
                    ph.print_info(
                        "All slices of stack '%s' were rejected entirely. "
                        "Information added." % stack.get_filename())

            ph.write_dictionary_to_json(
                deleted_slices_dic,
                os.path.join(
                    dir_output,
                    args.subfolder_motion_correction,
                    "rejected_slices.json"
                )
            )

    # ---------------------Final Volumetric Reconstruction---------------------
    ph.print_title("Final Volumetric Reconstruction")
    if args.sda:
        recon_method = sda.ScatteredDataApproximation(
            stacks,
            HR_volume,
            sigma=args.alpha,
            use_masks=args.use_masks_srr,
        )
    else:
        if args.reconstruction_type in ["TVL2", "HuberL2"]:
            recon_method = pd.PrimalDualSolver(
                stacks=stacks,
                reconstruction=HR_volume,
                reg_type="TV" if args.reconstruction_type == "TVL2" else "huber",
                iterations=args.iterations,
                use_masks=args.use_masks_srr,
            )
        else:
            recon_method = tk.TikhonovSolver(
                stacks=stacks,
                reconstruction=HR_volume,
                reg_type="TK1" if args.reconstruction_type == "TK1L2" else "TK0",
                use_masks=args.use_masks_srr,
            )
        recon_method.set_alpha(args.alpha)
        recon_method.set_iter_max(args.iter_max)
        recon_method.set_verbose(True)
    recon_method.run()
    time_reconstruction += recon_method.get_computational_time()
    HR_volume_final = recon_method.get_reconstruction()

    ph.print_subtitle("Final SDA Approximation Image Mask")
    SDA = sda.ScatteredDataApproximation(
        stacks, HR_volume_final, sigma=args.sigma, sda_mask=True)
    SDA.run()
    # HR volume contains updated mask based on SDA
    HR_volume_final = SDA.get_reconstruction()
    time_reconstruction += SDA.get_computational_time()

    elapsed_time_total = ph.stop_timing(time_start)

    # Write SRR result
    filename = recon_method.get_setting_specific_filename()
    HR_volume_final.set_filename(filename)
    dw.DataWriter.write_image(
        HR_volume_final.sitk,
        args.output,
        description=filename)
    dw.DataWriter.write_mask(
        HR_volume_final.sitk_mask,
        ph.append_to_filename(args.output, "_mask"),
        description=SDA.get_setting_specific_filename())

    HR_volume_iterations.insert(0, HR_volume_final)
    for stack in stacks:
        HR_volume_iterations.append(stack)

    if args.verbose:
        sitkh.show_stacks(
            HR_volume_iterations,
            segmentation=HR_volume_final,
            viewer=args.viewer,
        )

    # Summary
    ph.print_title("Summary")
    exe_file_info = os.path.basename(os.path.abspath(__file__)).split(".")[0]
    print("%s | Computational Time for Data Preprocessing: %s" %
          (exe_file_info, time_data_preprocessing))
    print("%s | Computational Time for Registrations: %s" %
          (exe_file_info, time_registration))
    print("%s | Computational Time for Reconstructions: %s" %
          (exe_file_info, time_reconstruction))
    print("%s | Computational Time for Entire Reconstruction Pipeline: %s" %
          (exe_file_info, elapsed_time_total))

    ph.print_line_separator()

    return 0
def main():

    # Read input
    input_parser = InputArgparser(
        description="Script to evaluate the similarity of simulated stack "
        "from obtained reconstruction against the original stack. "
        "This function takes the result of "
        "simulate_stacks_from_reconstruction.py as input.", )
    input_parser.add_filenames(required=True)
    input_parser.add_filenames_masks()
    input_parser.add_dir_output(required=True)
    input_parser.add_suffix_mask(default="_mask")
    input_parser.add_measures(default=["NCC", "SSIM"])
    input_parser.add_option(
        option_string="--prefix-simulated",
        type=str,
        help="Specify the prefix of the simulated stacks to distinguish them "
        "from the original data.",
        default="Simulated_",
    )
    input_parser.add_option(
        option_string="--dir-input-simulated",
        type=str,
        help="Specify the directory where the simulated stacks are. "
        "If not given, it is assumed that they are in the same directory "
        "as the original ones.",
        default=None)
    input_parser.add_slice_thicknesses(default=None)

    args = input_parser.parse_args()
    input_parser.print_arguments(args)

    # --------------------------------Read Data--------------------------------
    ph.print_title("Read Data")

    # Read original data
    filenames_original = args.filenames
    data_reader = dr.MultipleImagesReader(
        file_paths=filenames_original,
        file_paths_masks=args.filenames_masks,
        suffix_mask=args.suffix_mask,
        stacks_slice_thicknesses=args.slice_thicknesses,
    )
    data_reader.read_data()
    stacks_original = data_reader.get_data()

    # Read data simulated from obtained reconstruction
    if args.dir_input_simulated is None:
        dir_input_simulated = os.path.dirname(filenames_original[0])
    else:
        dir_input_simulated = args.dir_input_simulated
    filenames_simulated = [
        os.path.join("%s", "%s%s") %
        (dir_input_simulated, args.prefix_simulated, os.path.basename(f))
        for f in filenames_original
    ]
    data_reader = dr.MultipleImagesReader(filenames_simulated,
                                          suffix_mask=args.suffix_mask)
    data_reader.read_data()
    stacks_simulated = data_reader.get_data()

    for i in range(len(stacks_original)):
        try:
            stacks_original[i].sitk - stacks_simulated[i].sitk
        except:
            raise IOError(
                "Images '%s' and '%s' do not occupy the same space!" %
                (filenames_original[i], filenames_simulated[i]))

    similarity_measures = {
        m: SimilarityMeasures.similarity_measures[m]
        for m in args.measures
    }
    similarities = np.zeros(len(args.measures))

    for i in range(len(stacks_original)):
        nda_3D_original = sitk.GetArrayFromImage(stacks_original[i].sitk)
        nda_3D_simulated = sitk.GetArrayFromImage(stacks_simulated[i].sitk)
        nda_3D_mask = sitk.GetArrayFromImage(stacks_original[i].sitk_mask)

        path_to_file = os.path.join(
            args.dir_output,
            "Similarity_%s.txt" % stacks_original[i].get_filename())
        text = "# Similarity: %s vs %s (%s)." % (os.path.basename(
            filenames_original[i]), os.path.basename(
                filenames_simulated[i]), ph.get_time_stamp())
        text += "\n#\t" + ("\t").join(args.measures)
        text += "\n"
        ph.write_to_file(path_to_file, text, "w")
        for k in range(nda_3D_original.shape[0]):
            x_2D_original = nda_3D_original[k, :, :]
            x_2D_simulated = nda_3D_simulated[k, :, :]

            # zero slice, i.e. rejected during motion correction
            if np.abs(x_2D_simulated).sum() < 1e-6:
                x_2D_simulated[:] = np.nan
            x_2D_mask = nda_3D_mask[k, :, :]

            indices = np.where(x_2D_mask > 0)

            for m, measure in enumerate(args.measures):
                if len(indices[0]) > 0:
                    similarities[m] = similarity_measures[measure](
                        x_2D_original[indices], x_2D_simulated[indices])
                else:
                    similarities[m] = np.nan
            ph.write_array_to_file(path_to_file, similarities.reshape(1, -1))

    return 0