EHRTransferBenchmark/data_extraction_root.py at main · kingrc15/EHRTransferBenchmark · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
from __future__ import absolute_import
from __future__ import print_function

import os
import argparse
from data_extractor import utils

def data_extraction_root(args):
    output_dir = os.path.join(args.output_dir)
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    eicu_dir = os.path.join(args.eicu_dir)
    if not os.path.exists(eicu_dir):
        os.mkdir(eicu_dir)

    patients = utils.read_patients_table(args.eicu_dir, args.output_dir)
    stay_id = utils.cohort_stay_id(patients)
    utils.break_up_stays_by_unit_stay(patients, args.output_dir,stayid=stay_id, verbose=1)
    del patients

    print("reading lab table")
    lab = utils.read_lab_table(args.eicu_dir)
    utils.break_up_lab_by_unit_stay(lab, args.output_dir, stayid=stay_id, verbose=1)
    del lab

    print("reading nurseAssessment table")
    na = utils.read_na_table(args.eicu_dir)
    utils.break_up_na_by_unit_stay(na, args.output_dir, stayid=stay_id, verbose=1)
    del na

    print("reading nurseCharting table, might take some time")
    nc = utils.read_nc_table(args.eicu_dir)
    utils.break_up_stays_by_unit_stay_nc(nc, args.output_dir, stayid=stay_id, verbose=1)
    del nc

    # #Write the timeseries data into folders
    utils.extract_time_series_from_subject(args.output_dir, args.eicu_dir, args.task_dir)

    utils.delete_wo_timeseries(args.output_dir)

    #Write all the data into one dataframe
    utils.all_df_into_one_df(args.output_dir)


def main():
    parser = argparse.ArgumentParser(description="Create data for all tasks")
    parser.add_argument('--eicu_dir', type=str, help="Path to eICU dataset")
    parser.add_argument('--output_dir', type=str, help="Directory where the created data should be stored.")
    parser.add_argument('--task_dir', type=str, help="Directory where the task directories should be created")

    args, _ = parser.parse_known_args()

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    if not os.path.exists(args.eicu_dir):
        os.makedirs(args.eicu_dir)
    if not os.path.exists(args.task_dir):
        os.makedirs(args.task_dir)
        os.makedirs(os.path.join(args.task_dir, "root/"))
    data_extraction_root(args)
if __name__ == '__main__':
    main()