imageClassification/train_model.py at master · michelle14lee/imageClassification · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from sklearn import preprocessing
from sklearn import model_selection
from tensorflow import keras
from package import nn_model
from package import data_generator
import argparse
import csv
import pandas
import numpy
import os


def import_data(path = './data'):
    dataset = pandas.read_csv(os.path.join(path, 'data_labels.csv'), names = ['image_id', 'label'])
    dataset = dataset.to_numpy()
    numpy.random.shuffle(dataset)
    split_data = numpy.hsplit(dataset, 2)

    image_id = numpy.ravel(split_data[0])
    image_label, image_classes = encode_labels(numpy.ravel(split_data[1]))

    return image_id, image_label, image_classes

def encode_labels(label):
    label_encoder = preprocessing.LabelEncoder()
    label_encoder.fit(label)

    return label_encoder.transform(label), label_encoder.classes_


if __name__ == '__main__':
    image_id, image_label, image_classes = import_data()
    gparams = {
        'dim': (28, 28, 3),
        'batch_size': 64,
        'n_classes': len(image_classes),
    }

    X_train, X_test, y_train, y_test = model_selection.train_test_split(image_id, image_label, test_size = 0.2)

    image_train = data_generator.ImageGenerator(X_train, y_train, **gparams)
    image_val = data_generator.ImageGenerator(X_test, y_test, **gparams)

    image_model = nn_model.ImageClassModel(image_classes)
    image_model.fit(X_train = image_train, validation_data = image_val, epochs = 5)
    image_model.save_model(path = './model')