from tensorflow.keras.datasets import fashion_mnist

from tensorflow.keras.layers import Conv2D, Flatten, Dense, AvgPool2D

from tensorflow.keras.models import Sequential

from tensorflow.keras.optimizers import Adam

import numpy as np

 

 

def load_train(path):

    datagen = ImageDataGenerator(

        validation_split=0.25,

        rescale=1./255,

        horizontal_flip=True,

        vertical_flip=True)

        #rotation_range=90,

        #width_shift_range=0.2,

        #height_shift_range=0.2)  

   

    train_data = datagen.flow_from_directory(

        path,

        target_size=(150, 150),

        batch_size=16,

        class_mode='sparse',

        subset='training',

        seed=12345)

   

    return train_data

 

def create_model(input_shape):

    model = Sequential()

    model.add(Conv2D(filters=6, kernel_size=(5, 5), padding='same',

                 activation="relu", input_shape=input_shape))

    model.add(AvgPool2D(pool_size=(2, 2), strides=None, padding='valid'))

    model.add(Conv2D(filters=12, kernel_size=(5, 5), padding='valid',activation="relu"))

    model.add(Flatten())

    model.add(Dense(units=12, activation='softmax'))

    optimizer = Adam()

    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['acc'])

 

    return model

 

def train_model(model, train_data, test_data, batch_size=16, epochs=2,

                steps_per_epoch=None, validation_steps=None):

    model.fit(train_data,

              validation_data=test_data,

              batch_size=batch_size, epochs=epochs,

              steps_per_epoch=steps_per_epoch,

              validation_steps=validation_steps,

              verbose=2)

    return model