a_input = sio.loadmat('digits.mat')

train_labels = a_input['trainLabels'].T
test_labels = a_input['testLabels'].T

train_images = a_input['trainImages'].T
test_images = a_input['testImages'].T
train_columns = []
test_columns = []

# training_sample_size = len(train_images)
training_sample_size = int(sys.argv[1])
for i in range(training_sample_size):
    m = train_images[i][0]
    m = m / 255.0
    train_columns.append(np.ravel(m.T))
train_labels = train_labels[:training_sample_size]

test_sample_size = int(sys.argv[2])
for j in range(test_sample_size):
    m = test_images[j][0]
    m = m / 255.0
    test_columns.append(np.ravel(m.T))
test_labels = test_labels[:test_sample_size]

train_colArr = np.asarray(train_columns).T
test_colArr = np.asarray(test_columns).T