Model_SortToBids.py

import numpy as np
import os
import shutil
import yaml
import json

BIDS_VERSION = 'v1.2.1'


def create_directory(dir):
    if not os.path.exists(dir):
        os.makedirs(dir)


def get_out_file_name(subject_label, session_label, sub_type,
                      ce_label, modality_label, run_index):
    run_label = 'run-' + str(run_index)

    scan_labels = list(filter(None,
                             [subject_label, session_label,
                               ce_label, run_label, modality_label]))

    out_file_name = '_'.join(scan_labels) + '.nii.gz'

    out_directory = os.path.join(root_out_folder, subject_label,
                                 session_label, sub_type)

    out_file = os.path.join(out_directory, out_file_name)

    return out_file, out_directory

def stb(nifti_dir):
    for root, dir, files in os.walk(nifti_dir):
        if len(files) > 0:
            patient_ID = os.path.basename(os.path.normpath(root))
            patient_ID = patient_ID.replace('-', '')
            subject_label = 'sub-' + patient_ID
            session_dict = dict()
            cur_session_index = 1
            for i_file in files:
                print(i_file)
                # Reset everything
                sub_type = None
                modality_label = None
                ce_label = None
                session_label = None

                full_file = os.path.join(root, i_file)
                i_file = i_file.split('.nii.gz')[0]
                splitted_elements = i_file.split('__')
                session_ID = splitted_elements[2]

                if session_ID not in session_dict:
                    session_dict[session_ID] = cur_session_index
                    session_num = cur_session_index
                    cur_session_index += 1
                else:
                    session_num = session_dict[session_ID]
                session_label = 'ses-' + str(session_num)

                indices = np.argwhere([i_file in i_result_name for i_result_name in file_names])
                if len(indices) != 0:
                    predictions = prediction_results[indices].ravel()
                    i_prediction = np.bincount(predictions).argmax() - 1
                else:
                    i_prediction = -1

                if i_prediction == 0 or i_prediction == 1:
                    sub_type = 'anat'
                    modality_label = 'T1w'
                    if i_prediction == 1:
                        ce_label = 'ce-GD'
                elif i_prediction == 2:
                    sub_type = 'anat'
                    modality_label = 'T2w'
                elif i_prediction == 3:
                    sub_type = 'anat'
                    modality_label = 'PD'
                elif i_prediction == 4:
                    sub_type = 'anat'
                    modality_label = 'FLAIR'
                elif i_prediction == 5:
                    sub_type = 'dwi'
                    modality_label = 'dwi'
                else:
                    break
                current_run_index = 1

                scan_labels = [subject_label, ce_label, modality_label]

                out_file, out_directory = get_out_file_name(subject_label, session_label,
                                             sub_type, ce_label, modality_label,
                                             current_run_index)

                while os.path.exists(out_file):
                    current_run_index += 1
                    out_file, out_directory = get_out_file_name(subject_label, session_label,
                                                 sub_type, ce_label, modality_label,
                                                 current_run_index)

                if not os.path.exists(out_directory):
                    os.makedirs(out_directory)

                shutil.copy(full_file, out_file)


def main(configFile):
    with open(configFile, 'r') as ymlfile:
        cfg = yaml.load(ymlfile)

    prediction_file = cfg['post_processing']['prediction_file']
    root_dicom_folder = cfg['preprocessing']['root_dicom_folder']
    base_dir = os.path.dirname(os.path.normpath(root_dicom_folder))
    nifti_dir = os.path.join(base_dir, 'NIFTI')

    base_dir = os.path.dirname(os.path.normpath(root_dicom_folder))
    root_out_folder = os.path.join(base_dir, 'BIDS_SORTED')

    os.makedirs(root_out_folder, exist_ok=True)

    predictions = np.loadtxt(prediction_file, dtype=np.str)

    prediction_names = ['T1', 'T1GD', 'T2', 'PD', 'FLAIR', 'DWI_DWI', 'DERIVED', 'PWI_DSC', 'UNKNOWN']
    orientation_names = ['3D', 'Ax', 'Cor', 'Sag', 'Obl', '4D', 'UNKNOWN']

    prediction_file_names = predictions[:, 0]
    prediction_results = predictions[:, 1].astype(np.int)

    file_names = [i_file_name.split(os.sep)[-1] for i_file_name in prediction_file_names]

    unique_names = np.unique(file_names)
    unique_predictions = np.zeros([len(unique_names), 1])

    out_json = {'Name': 'DDS_sorted_dataset',
                'BIDSVersion': BIDS_VERSION}

    with open(os.path.join(root_out_folder, 'dataset_description.json'), 'w') as jf:
        json.dump(out_json, jf)

    stb(nifti_dir)