mines_vs_rocks/src/featurization.py

#######################################################
# Receives train and test *.csv sets and generates features for each of them
# Returns train and test  *.joblib sets with new features

import pandas as pd
import os
import errno
import sys
import joblib


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 get_df(path):
    df = pd.read_csv(path)
    sys.stderr.write('The input data frame {} size is {}\n'.format(path, df.shape))
    return df

def featurize(train_path, test_path):
    """
    Generates features for the train and test sets
    # TODO: for now it's only the identity function to stage the pipeline
    :param train_path: path to input train dataset, csv file
    :param test_path: path to input test dataset, csv file
    :return: tuple, (train, test), featurized pandas dataframes
    """
    df_train = get_df(train_path)
    features_train = df_train

    df_test = get_df(test_path)
    features_test = df_test
    return features_train, features_test

if __name__ == "__main__":
    if len(sys.argv) != 3 and len(sys.argv) != 5:
        sys.stderr.write('Arguments error. Usage:\n')
        sys.stderr.write('\tpython featurization.py data-dir-path features-dir-path\n')
        sys.exit(1)

    mkdir_p(sys.argv[2])

    train_input = os.path.join(sys.argv[1], 'train.joblib')
    test_input = os.path.join(sys.argv[1], 'test.joblib')
    train_output = os.path.join(sys.argv[2], 'train.joblib')
    test_output = os.path.join(sys.argv[2], 'test.joblib')

    # try:
    #     reload(sys)
    #     sys.setdefaultencoding('utf-8')
    # except NameError:
    #     pass


    train_features, test_features = featurize(train_input, test_input)
    joblib.dump(train_features, train_output)
    joblib.dump(test_features, test_output)