mines_vs_rocks/src/preprocessing.py

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import os
import errno
import sys

OUTPUT = os.path.join('data', 'preprocessed')
split = 0.2
random_seed = 42


def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def preprocess(path):
    input = pd.read_csv(path, header=None)
    return(input)

def prepare_dataset(df):
    df.rename({60: 'label'}, axis='columns', inplace=True)
    df['label'] = pd.Series(LabelEncoder().fit_transform(df['label'])).astype('category')
    # y = df['label'].copy(deep=True)
    # label_enc = LabelEncoder().fit(y)
    # y = label_enc.transform(y)

    # X = df.drop('label', axis=1).astype(float)
    # X_train, X_test, y_train, y_test = train_test_split(X.values, y, test_size=split,
    #                                                     shuffle=True, random_state=random_seed)
    # df_train = pd.DataFrame(X_train).join(pd.Series(y_train, name='label'))
    # df_test = pd.DataFrame(X_test).join(pd.Series(y_test, name='label'))
    df_train, df_test = train_test_split(df, stratify=df['label'], test_size=split, shuffle=True, random_state=random_seed)
    return df_train, df_test


if __name__ == "__main__":
    if len(sys.argv) != 2:
        sys.stderr.write('Arguments error. Usage:\n')
        sys.stderr.write('\tpython prepare.py data\n')
        sys.exit(1)

    mkdir_p(OUTPUT)
    input = preprocess(sys.argv[1])
    train, test = prepare_dataset(input)
    train.to_csv(os.path.join(OUTPUT, 'train.joblib'), index_label=False)
    test.to_csv(os.path.join(OUTPUT, 'test.joblib'), index_label=False)
    print('Saved train and test sets in ' + OUTPUT)