Function prepare_datasets

train.py:265–302

(args)

    writer.close()

def prepare_datasets(args):

    df = pd.read_csv(args.manifest)

    n_bins = len(df['disc_label'].unique())
    assert n_bins == args.n_bins, 'mismatch between the number of bins passed in args and classes in dataset'
    bins_values = get_bins_time_value(df, n_bins, time_col_name='recurrence_years', label_time_col_name='disc_label')
    assert len(bins_values)==n_bins
    print(f'Read {args.manifest} dataset containing {len(df)} samples with {n_bins} bins of following values {bins_values}')

    # NOTE: you may need to use the two lines below depending on how the category is listed in the csv file. 
    #df.stage = df.stage.apply(lambda x : 'III' if 'III' in x else ('II' if 'II' in x else 'I')).astype("category")
    #df.stage = pd.Categorical(df['stage'], categories=['I', 'II', 'III'], ordered=True).codes
    print(f'stage taxonomy used: {df.stage.unique()}')

    try:
        training_set = df[df["split"] == "training"]
        validation_set = df[df["split"] == "validation"]
    except:
        raise Exception(
            f"Could not find training and validation splits in {args.manifest}"
        )

    train_split = FeatureBagsDataset(df=training_set,
                                    data_dir=args.data_dir,
                                    input_feature_size=args.input_feature_size, 
                                    stage_class=len(training_set.stage.unique()))

    val_split = FeatureBagsDataset(df=validation_set,
                                    data_dir=args.data_dir, 
                                    input_feature_size=args.input_feature_size, 
                                    stage_class=len(validation_set.stage.unique()))

    # To compute the Brier score (BS), you need a specific format of censorship and times.
    _, train_BS = get_survival_data_for_BS(training_set, time_col_name='recurrence_years')
    _, test_BS = get_survival_data_for_BS(validation_set, time_col_name='recurrence_years')

    return train_split, val_split, train_BS, test_BS, bins_values, len(df.stage.unique())


def main(args):