data_preprocessing.py

# import pickle
#
# path_root = '/home/xinyue/VQA_ReGat/data/mimic_vqa/'
# train_dataset_path = path_root + 'mimic_dataset_train.pkl'
# val_dataset_path = path_root + 'mimic_dataset_val.pkl'
# test_dataset_path = path_root + 'mimic_dataset_test.pkl'
# train_dataset = pickle.load(open(train_dataset_path, 'rb'))
# val_dataset = pickle.load(open(val_dataset_path, 'rb'))
# test_dataset = pickle.load(open(test_dataset_path, 'rb'))
# total_dataset = train_dataset + val_dataset + test_dataset



import argparse
import pickle
import h5py
import pandas as pd
from tqdm import tqdm
import random
import numpy as np
import os
import glob
import json
# needed files
# dictionary file: mimic_dictionary_full.pkl
# ans2label file: mimic_ans2label_full.pkl
# label2ans file: mimic_label2ans_full.pkl
# features file: cmb_bbox_features_full.hdf5
# split files: mimic_dataset_train_full.pkl and mimic_dataset_val_full.pkl and mimic_dataset_test_full.pkl
# the most important file is the split file

def find_name_id_in_dd_report_name(name, d_d):
    for i in range(len(d_d)):
        names = d_d.iloc[i]['report_name'].split(';')
        for n in names:
            if n in name:
                official_name = d_d.iloc[i]['official_name']
                return official_name, d_d.iloc[i]['id']
    return None

def check_existing_files():
    dictionary_path = 'data/medical_cxr_vqa/mimic_dictionary.pkl'
    ans2label_path =  'data/medical_cxr_vqa/mimic_ans2label.pkl'
    label2ans_path = 'data/medical_cxr_vqa/mimic_label2ans.pkl'
    features_path = 'data/medical_cxr_vqa/cmb_bbox_features.hdf5'
    train_path = 'data/medical_cxr_vqa/mimic_dataset_train.pkl'
    val_path = 'data/medical_cxr_vqa/mimic_dataset_val.pkl'
    test_path = 'data/medical_cxr_vqa/mimic_dataset_test.pkl'

    # read dictionary
    with open(dictionary_path, 'rb') as f:
        dictionary = pickle.load(f)
    # read ans2label
    with open(ans2label_path, 'rb') as f:
        ans2label = pickle.load(f)
    # read label2ans
    with open(label2ans_path, 'rb') as f:
        label2ans = pickle.load(f)
    # read from hdft5 file
    hf = h5py.File(features_path, 'r')
    # get length of the features

    # read train
    with open(train_path, 'rb') as f:
        train = pickle.load(f)
    with open(val_path, 'rb') as f:
        val = pickle.load(f)
    print('a')


def assign_splits(record, ans_split_set,train_dataset, val_dataset, test_dataset,label2ans):
    for ans in record['answer']['answer']:
        if ans not in ans_split_set['test']:
            ans_split_set['test'].add(ans)
            test_dataset.append(record)
            # add the rest of answers to ans_split_set['test']
            for ans in record['answer']['answer']:
                ans_split_set['test'].add(ans)
            return ans_split_set, train_dataset, val_dataset, test_dataset

        elif ans not in ans_split_set['val']:
            ans_split_set['val'].add(ans)
            val_dataset.append(record)
            # add the rest of answers to ans_split_set['val']
            for ans in record['answer']['answer']:
                ans_split_set['val'].add(ans)
            return ans_split_set, train_dataset, val_dataset, test_dataset
        elif ans not in ans_split_set['train']:
            ans_split_set['train'].add(ans)
            train_dataset.append(record)
            # add the rest of answers to ans_split_set['train']
            for ans in record['answer']['answer']:
                ans_split_set['train'].add(ans)
            return ans_split_set, train_dataset, val_dataset, test_dataset
    # when random number is less than 0.8
    if random.random() < 0.8:
        train_dataset.append(record)
    else:
        if random.random() < 0.5:
            val_dataset.append(record)
        else:
            test_dataset.append(record)
    # elif random.random() < 0.9:
    #     val_dataset.append(record)
    # else:
    #     test_dataset.append(record)
    return ans_split_set, train_dataset, val_dataset, test_dataset


def preprocess_dataset(dataroot='data/medical_cxr_vqa/', remove_tail = False, less_yes_no = True, filter_low_freq = False):
    path = os.path.join(dataroot, 'medical-cxr-vqa-questions.csv')
    # read csv file using pandas
    df = pd.read_csv(path)
    mimic_all_path = os.path.join(dataroot, 'mimic_all.csv')
    d_all = pd.read_csv(mimic_all_path)
    mimic_shape_path = os.path.join(dataroot, 'mimic_shape_full.pkl')
    with open(mimic_shape_path, 'rb') as f:
        mimic_shape = pickle.load(f)
    mimic_shapeid_path = os.path.join(dataroot, 'mimic_shapeid_full.pkl')
    with open(mimic_shapeid_path, 'rb') as f:
        mimic_shapeid = pickle.load(f)

    wordset = set()
    answerset = set()
    answer_count = {}
    # obtain labelse first
    for i in tqdm(range(len(df))):
        if df.iloc[i]['question_type'] != 'difference':
            question = df.iloc[i]['question'].replace('?',' ?')
            answers = df.iloc[i]['answer'].replace('.','')
            answers = answers.split(', ')
            # if len(answers) > 1:
            #     print('a')
            for answer in answers:
                if answer not in answer_count:
                    answer_count[answer] = 1
                else:
                    answer_count[answer] += 1
            wordset.update(question.split())
            wordset.update(answers)
        # if i >= 100000:
        #     break
    wordset = list(wordset)
    if remove_tail:
        # remove answers that count is less than 5 from ans2label
        label2ans = [label for label in answer_count if answer_count[label] >= 5]
    else:
        label2ans = list(answer_count.keys())
    # wordset.sort()
    # label2ans.sort()
    # get word2id
    word2id = {word: i for i, word in enumerate(wordset)}
    # transform labelset to dict
    ans2label = {label: i for i, label in enumerate(label2ans)}
    answerset = set(label2ans)


    total_dataset = []
    for i in tqdm(range(len(df))):
        if df.iloc[i]['question_type'] != 'difference':
            record = {}
            question = df.iloc[i]['question']
            answer = df.iloc[i]['answer'].replace('.', '')
            answer = answer.split(', ')
            anss = answer.copy()
            while anss:
                ans = anss.pop()
                if ans not in answerset:
                    # remove ans form answer
                    answer.remove(ans)
            if answer == []:
                continue
            subject_id = df.iloc[i]['subject_id']
            study_id = df.iloc[i]['study_id']
            # find dicom_id from d_all and ('view' is 'postero-anterior' or 'antero-posterior')
            dicom_id = d_all[d_all['study_id'] == study_id]
            dicom_id = dicom_id[dicom_id['view'].isin(['postero-anterior','antero-posterior'])]['dicom_id'].values[0]
            # dicom_id = dicom_id
            # get labels
            labels = [ans2label[ans] for ans in answer]
            # set scores to all 1.0
            scores = [1.0] * len(labels)
            height = mimic_shape[mimic_shapeid[dicom_id]]['height']
            width = mimic_shape[mimic_shapeid[dicom_id]]['width']
            image = mimic_shapeid[dicom_id]

            record['subject_id'] = subject_id
            record['study_id'] = study_id
            record['dicom_id'] = dicom_id
            record['question'] = question
            record['question_type'] = df.iloc[i]['question_type']
            record['answer'] = {'labels': labels, 'scores': scores, 'answer': answer}
            record['height'] = height
            record['width'] = width
            record['image'] = image
            total_dataset.append(record)
        # if i >= 100000:
        #     break
    dictionary = [word2id, wordset]
    # split dataset to train, val and test
    train_dataset = []
    val_dataset = []
    test_dataset = []
    ans_split_set = {'train': set(), 'val': set(), 'test': set()}
    # split the datasets
    #1
    # for i in range(len(total_dataset)):
    #     ans_split_set, train_dataset, val_dataset, test_dataset = assign_splits(total_dataset[i], ans_split_set, train_dataset, val_dataset, test_dataset,label2ans)
    #2 avoid the same patient occurs in train, val and test
    train_dataset = total_dataset[:int(len(total_dataset) * 0.8)]
    val_dataset = total_dataset[int(len(total_dataset) * 0.8):int(len(total_dataset) * 0.9)]
    test_dataset = total_dataset[int(len(total_dataset) * 0.9):]
    print('train:', len(train_dataset))
    print('val:', len(val_dataset))
    print('test:', len(test_dataset))
    # save to pickle file
    total_dataset_path = os.path.join(dataroot, 'total_dataset.pkl')
    train_path = os.path.join(dataroot, 'mimic_dataset_train.pkl')
    val_path = os.path.join(dataroot, 'mimic_dataset_val.pkl')
    test_path = os.path.join(dataroot, 'mimic_dataset_test.pkl')
    with open(total_dataset_path, 'wb') as f:
        pickle.dump(total_dataset, f)
    with open(train_path, 'wb') as f:
        pickle.dump(train_dataset, f)
    with open(val_path, 'wb') as f:
        pickle.dump(val_dataset, f)
    with open(test_path, 'wb') as f:
        pickle.dump(test_dataset, f)
    #save dictionary
    dictionary_path = os.path.join(dataroot, 'mimic_dictionary.pkl')
    with open(dictionary_path, 'wb') as f:
        pickle.dump(dictionary, f)
    # save label2ans
    label2ans_path = os.path.join(dataroot, 'mimic_label2ans.pkl')
    with open(label2ans_path, 'wb') as f:
        pickle.dump(label2ans, f)
    # save ans2label
    ans2label_path = os.path.join(dataroot, 'mimic_ans2label.pkl')
    with open(ans2label_path, 'wb') as f:
        pickle.dump(ans2label, f)

    # save label_count
    label_count_path = os.path.join(dataroot, 'mimic_label_count.pkl')
    with open(label_count_path, 'wb') as f:
        pickle.dump(answer_count, f)



def get_answerset(split):
    answerset = set()
    for i in range(len(split)):
        for ans in split[i]['answer']['answer']:
            answerset.add(ans)
    return answerset

def read_file(dataroot='data/medical_cxr_vqa/'):
    train_path = os.path.join(dataroot, 'mimic_dataset_train.pkl')
    val_path = os.path.join(dataroot, 'mimic_dataset_val.pkl')
    test_path = os.path.join(dataroot, 'mimic_dataset_test.pkl')
    ans2label_path = os.path.join(dataroot, 'mimic_ans2label.pkl')
    label2ans_path = os.path.join(dataroot, 'mimic_label2ans.pkl')
    lable_count_path = os.path.join(dataroot, 'mimic_label_count.pkl')
    dictionary_path = os.path.join(dataroot, 'mimic_dictionary.pkl')
    train = pickle.load(open(train_path, 'rb'))
    val = pickle.load(open(val_path, 'rb'))
    test = pickle.load(open(test_path, 'rb'))
    ans2label = pickle.load(open(ans2label_path, 'rb'))
    label2ans = pickle.load(open(label2ans_path, 'rb'))
    label_count = pickle.load(open(lable_count_path, 'rb'))
    train_ans = get_answerset(train)
    val_ans = get_answerset(val)
    test_ans = get_answerset(test)
    dictionary = pickle.load(open(dictionary_path, 'rb'))

    count = {'abnormality': 0, 'view': 0, 'presence': 0, 'location': 0, 'level':0, 'type':0}
    for split in [train, val, test]:
        for i in range(len(split)):
            ques_type = split[i]['question_type']
            count[ques_type] += 1
    # calculate the percentage of each question type
    for key in count:
        print(key,count[key],  count[key]/(len(train)+len(val)+len(test)))
    print('total:', len(train)+len(val)+len(test))

    idset = set()
    for split in [train, val, test]:
        for i in range(len(split)):
            idset.add(split[i]['study_id'])
    print('idset:', len(idset))

    idcount = {}
    for split in [train, val, test]:
        for i in range(len(split)):
            id = split[i]['study_id']
            if id not in idcount:
                idcount[id] = 1
            else:
                idcount[id] += 1
    countidcount = {}
    for id in idcount:
        countidcount[idcount[id]] = countidcount.get(idcount[id], 0) + 1
    print('countidcount:', countidcount)

    answerset = set()
    for split in [train, val, test]:
        for i in range(len(split)):
            for ans in split[i]['answer']['answer']:
                answerset.add(ans)
    print('answerset:', len(answerset))

def resplit_dataset(dataroot='data/medical_cxr_vqa/'):
    # read_file()
    total_dataset_path = os.path.join(dataroot, 'total_dataset.pkl')
    label2ans_path = os.path.join(dataroot, 'mimic_label2ans.pkl')
    total_dataset = pickle.load(open(total_dataset_path, 'rb'))
    label2ans = pickle.load(open(label2ans_path, 'rb'))
    ans_split_set = {'train': set(), 'val': set(), 'test': set()}
    train_dataset = []
    val_dataset = []
    test_dataset = []

    for i in range(len(total_dataset)):
        ans_split_set, train_dataset, val_dataset, test_dataset = assign_splits(total_dataset[i], ans_split_set, train_dataset, val_dataset, test_dataset,label2ans)

    print('train:', len(train_dataset))
    print('val:', len(val_dataset))
    print('test:', len(test_dataset))
    # save to pickle file
    train_path = os.path.join(dataroot, 'mimic_dataset_train.pkl')
    val_path = os.path.join(dataroot, 'mimic_dataset_val.pkl')
    test_path = os.path.join(dataroot, 'mimic_dataset_test.pkl')
    with open(total_dataset_path, 'wb') as f:
        pickle.dump(total_dataset, f)
    with open(train_path, 'wb') as f:
        pickle.dump(train_dataset, f)
    with open(val_path, 'wb') as f:
        pickle.dump(val_dataset, f)
    with open(test_path, 'wb') as f:
        pickle.dump(test_dataset, f)

def remove_low_freq_labels(dataroot='data/medical_cxr_vqa/'):
    train_path = os.path.join(dataroot, 'mimic_dataset_train.pkl')
    val_path = os.path.join(dataroot, 'mimic_dataset_val.pkl')
    test_path = os.path.join(dataroot, 'mimic_dataset_test.pkl')
    label2ans_path = os.path.join(dataroot, 'mimic_label2ans.pkl')
    ans2label_path = os.path.join(dataroot, 'mimic_ans2label.pkl')
    label_count_path = os.path.join(dataroot, 'mimic_label_count.pkl')

    train = pickle.load(open(train_path, 'rb'))
    val = pickle.load(open(val_path, 'rb'))
    test = pickle.load(open(test_path, 'rb'))
    label2ans = pickle.load(open(label2ans_path, 'rb'))
    ans2label = pickle.load(open(ans2label_path, 'rb'))
    label_count = pickle.load(open(label_count_path, 'rb'))

    labels_need_to_remove_total = []
    for split in [train, val, test]:
        labels_need_to_remove = label2ans.copy() + label2ans.copy()
        for i in tqdm(range(len(split))):
            # for label in split[i]['answer']['labels']:
            #     if label == 4:
            #         print('label 4')
            for ans in split[i]['answer']['answer']:
                if ans == 'apical right area':
                    print('a')
                try:
                    labels_need_to_remove.remove(ans)
                except:
                    pass
        labels_need_to_remove_total += labels_need_to_remove
    labels_need_to_remove_total = set(labels_need_to_remove_total)
    print('total number of labels need to remove:', len(labels_need_to_remove_total))
    for ans in labels_need_to_remove_total:
        print(ans, label_count[ans])

    # remove labels from label2ans
    answers = label2ans.copy()
    while answers:
        ans = answers.pop()
        if ans in labels_need_to_remove_total:
            label2ans.remove(ans)
    ans2label = {ans: i for i, ans in enumerate(label2ans)}

    # remove labels from splits
    splits = [train, val, test]
    for k, split in enumerate(splits):
        mask = np.ones(len(split), dtype=bool)
        for i in tqdm(range(len(split))):
            # if i == 8621:
            #     print('a')
            split[i]['answer']['labels'] = []
            # remove low freq labels ans answers
            answers = split[i]['answer']['answer'].copy()
            while answers:
                ans = answers.pop()
            # for ans in split[i]['answer']['answer']:
                if ans in labels_need_to_remove_total:
                    split[i]['answer']['answer'].remove(ans)
                    # split[i]['answer']['labels'].remove(ans2label[ans])
                    split[i]['answer']['scores'].pop()
                    if split[i]['answer']['answer'] == []:
                        # remove this record
                        mask[i] = False
            # reassign new labels
            for ans in split[i]['answer']['answer']:
                try:
                    split[i]['answer']['labels'].append(ans2label[ans])
                except:
                    print('a')
        # sample list by mask
        splits[k] = [split[i] for i in range(len(split)) if mask[i]]





    # reassign labels to splits
    for k, split in enumerate(splits):
        for i in range(len(split)):
            labels = []
            for ans in split[i]['answer']['answer']:
                labels.append(ans2label[ans])
            split[i]['answer']['labels'] = labels

    train, val, test = splits
    print('train:', len(train))
    print('val:', len(val))
    print('test:', len(test))

    # save files
    with open(label2ans_path, 'wb') as f:
        pickle.dump(label2ans, f)
    with open(ans2label_path, 'wb') as f:
        pickle.dump(ans2label, f)
    with open(train_path, 'wb') as f:
        pickle.dump(train, f)
    with open(val_path, 'wb') as f:
        pickle.dump(val, f)
    with open(test_path, 'wb') as f:
        pickle.dump(test, f)



def less_yes_no(dataroot='data/medical_cxr_vqa/'):
    train_path = os.path.join(dataroot, 'mimic_dataset_train.pkl')
    val_path = os.path.join(dataroot, 'mimic_dataset_val.pkl')
    test_path = os.path.join(dataroot, 'mimic_dataset_test.pkl')

    train = pickle.load(open(train_path, 'rb'))
    val = pickle.load(open(val_path, 'rb'))
    test = pickle.load(open(test_path, 'rb'))

    print('train:', len(train))
    print('val:', len(val))
    print('test:', len(test))

    new_train = []
    new_val = []
    new_test = []
    new_splits = [new_train, new_val, new_test]
    splits = [train, val, test]
    answer_count = {}
    for k, split in enumerate(splits):
        for i in range(len(split)):
            if split[i]['answer']['answer'] == ['yes'] or split[i]['answer']['answer'] == ['no']:
                rand = random.random()
                if rand < 0.1:
                    new_splits[k].append(split[i])
                    answer = split[i]['answer']['answer'][0]
                    if answer not in answer_count:
                        answer_count[answer] = 1
                    else:
                        answer_count[answer] += 1
            else:
                new_splits[k].append(split[i])
                answers = split[i]['answer']['answer']
                for answer in answers:
                    if answer not in answer_count:
                        answer_count[answer] = 1
                    else:
                        answer_count[answer] += 1
    train, val, test = new_splits
    print('train:', len(train))
    print('val:', len(val))
    print('test:', len(test))

    # save files
    with open(train_path, 'wb') as f:
        pickle.dump(train, f)
    with open(val_path, 'wb') as f:
        pickle.dump(val, f)
    with open(test_path, 'wb') as f:
        pickle.dump(test, f)

    label_count_path = os.path.join(dataroot, 'mimic_label_count.pkl')
    with open(label_count_path, 'wb') as f:
        pickle.dump(answer_count, f)

def check_statistics(path_root = 'data/medical_cxr_vqa/'):
    train_dataset_path = path_root + 'mimic_dataset_train.pkl'
    val_dataset_path = path_root + 'mimic_dataset_val.pkl'
    test_dataset_path = path_root + 'mimic_dataset_test.pkl'
    train_dataset_path = pickle.load(open(train_dataset_path, 'rb'))
    val_dataset_path = pickle.load(open(val_dataset_path, 'rb'))
    test_dataset_path = pickle.load(open(test_dataset_path, 'rb'))
    total_dataset = train_dataset_path + val_dataset_path + test_dataset_path

    abn = 0
    pres = 0
    view = 0
    type = 0
    level = 0
    loc = 0
    yesno = 0
    dicom_set = set()
    for i in range(len(total_dataset)):
        question_type = total_dataset[i]['question_type']
        answer = total_dataset[i]['answer']['answer']
        dicom = total_dataset[i]['dicom_id']
        dicom_set.add(dicom)
        if answer[0] == 'yes' or answer[0] == 'no':
            yesno += 1
        if question_type == 'abnormality':
            abn += 1
        elif question_type == 'presence':
            pres += 1
        elif question_type == 'view':
            view += 1
        elif question_type == 'type':
            type += 1
        elif question_type == 'level':
            level += 1
        elif question_type == 'location':
            loc += 1
    total = [abn, pres, view, type, level, loc]
    print('abn, pres, view, type, level, loc:', total)
    print('yesno:', yesno)
    t = 0
    for val in total:
        t += val
    print('total number:', t)
    print('dicom number:', len(dicom_set))


    answer_count = {}
    for i in range(len(total_dataset)):
        label = total_dataset[i]['answer']['answer']
        for l in label:
            if l not in answer_count:
                answer_count[l] = 1
            else:
                answer_count[l] += 1
    print('answer count:', answer_count)

    # sort answer count
    # dict to tuple
    answer_count_tuple = []
    for key, value in answer_count.items():
        answer_count_tuple.append((key, value))
    answer_count_tuple.sort(key=lambda x: x[1], reverse=True)
    print('answer count:', answer_count_tuple)


    answer_count_list = []
    for key in answer_count:
        answer_count_list.append(answer_count[key])
    # answer_count_list.sort()
    print('answer count list:', answer_count_list)
    print('total number of answers:', len(answer_count_list))
    return answer_count_tuple

def get_disease_graph_node_labels(disease_file_path='/home/xinyue/chatgpt/output/all_diseases_standardized4.json', dataroot='data/medical_cxr_vqa/'):
    disease_file = json.load(open(disease_file_path, 'r'))

    disease_lib_path = 'lib/disease_lib_llm_full.csv'
    disease_lib = pd.read_csv(disease_lib_path)
    disease_names = disease_lib['official_name'].tolist()
    name2id = {}
    for i, name in enumerate(disease_names):
        name2id[name] = i

    node_labels = {}




    for record in tqdm(disease_file):
        study_id = record['study_id']
        label = []
        for ent in record['entity']:
            if 'probability_score' not in record['entity'][ent] or record['entity'][ent]['probability_score'] > 0:
                # if ent not in node_label2id:
                #     node_label2id[ent] = len(node_label2id)
                ret = find_name_id_in_dd_report_name(ent, disease_lib)
                name = ret[0]
                try:
                    label.append(name2id[name])
                except:
                    print(ent, 'not added. could because its frequency is too low')
        node_labels[study_id] = label

    # save
    node_label_path = os.path.join(dataroot, 'node_labels.pkl')
    with open(node_label_path, 'wb') as f:
        pickle.dump(node_labels, f)

    # node_label2id_path = '../data/mimic_vqa/node_label2id.pkl'
    # with open(node_label2id_path, 'wb') as f:
    #     pickle.dump(node_label2id, f)
    train_dataset_path = os.path.join(dataroot, 'mimic_dataset_train.pkl')
    val_dataset_path = os.path.join(dataroot, 'mimic_dataset_val.pkl')
    test_dataset_path = os.path.join(dataroot, 'mimic_dataset_test.pkl')
    train_dataset_path = pickle.load(open(train_dataset_path, 'rb'))
    val_dataset_path = pickle.load(open(val_dataset_path, 'rb'))
    test_dataset_path = pickle.load(open(test_dataset_path, 'rb'))
    total_dataset = train_dataset_path + val_dataset_path + test_dataset_path

def any_in_list(a, b):
    for i in a:
        if i in b:
            return True
    return False




def preprocess_csv2pkl(dataroot='data/medical_cxr_vqa/'):
    path = os.path.join(dataroot, 'medical-cxr-vqa-questions.csv')
    df = pd.read_csv(path)

    mimic_shape_path = 'data/mimic_shape_full.pkl'
    with open(mimic_shape_path, 'rb') as f:
        mimic_shape = pickle.load(f)
    mimic_shapeid_path = 'data/mimic_shapeid_full.pkl'
    with open(mimic_shapeid_path, 'rb') as f:
        mimic_shapeid = pickle.load(f)
    ans2label_path = os.path.join(dataroot,'mimic_ans2label.pkl')
    with open(ans2label_path, 'rb') as f:
        ans2label = pickle.load(f)

    train_set = []
    val_set = []
    test_set = []
    for i in tqdm(range(len(df))):
        record = df.iloc[i]
        # tranform record to dict
        record_dict = record.to_dict()
        image = mimic_shapeid[record_dict['dicom_id']]
        height = mimic_shape[image]['height']
        width = mimic_shape[image]['width']
        split = record_dict['split']

        answer = df.iloc[i]['answer'].replace('.', '')
        answer = answer.split(',')
        # labels = [ans2label[ans] for ans in answer]
        labels = []
        cache = ''
        for ans in answer:
            if ans in ans2label:
                labels.append(ans2label[ans])
            elif cache in ans2label:
                labels.append(ans2label[cache])
                cache = ''
            else:
                cache += ans
        scores = [1.0] * len(labels)

        record_dict['answer'] = {'labels': labels, 'scores': scores, 'answer': answer}
        record_dict['height'] = height
        record_dict['width'] = width
        record_dict['image'] = image

        if split =='train':
            train_set.append(record_dict['dicom_id'])
        elif split == 'val':
            val_set.append(record_dict['dicom_id'])
        elif split == 'test':
            test_set.append(record_dict['dicom_id'])

    train_path = os.path.join(dataroot,'mimic_dataset_train.pkl')
    val_path = os.path.join(dataroot,'mimic_dataset_val.pkl')
    test_path = os.path.join(dataroot,'mimic_dataset_test.pkl')

    with open(train_path, 'wb') as f:
        pickle.dump(train_set, f)
    with open(val_path, 'wb') as f:
        pickle.dump(val_set, f)
    with open(test_path, 'wb') as f:
        pickle.dump(test_set, f)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', type=str, default='simple', choices=['simple', 'full'], help='simple is for using the provided csv dataset directly. Full is for preprocessing all the way from the LLM generated dataset')

    args = parser.parse_args()

    return args

if __name__ == '__main__':
    args = parse_args()

    if args.mode == 'simple':
        preprocess_csv2pkl()
    elif args.model == 'full':
        preprocess_dataset(remove_tail=False, less_yes_no=False, filter_low_freq=True)
        remove_low_freq_labels()
        get_disease_graph_node_labels()