from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.datasets import mnist
from keras.layers import Activation, Dropout, Flatten, Dense
from keras.layers import Conv2D, MaxPooling2D
from keras import backend as K
import keras
import numpy
import os

save_dir = os.path.join(os.getcwd(), 'catsvsdogs_saved')

train_folder = 'train'
test_folder = 'test'
img_w, img_h = 150,150
input_shape = (img_w, img_h, 3)
epochs = 15
batch_size = 20
train_samples = 24800
test_samples = 200

model = Sequential()

model.add(Conv2D(32, (3,3), input_shape=input_shape))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Conv2D(32, (3,3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Conv2D(64, (3,3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Flatten())
model.add(Dense(64))
model.add(Activation('relu'))
model.add(Dropout(0.5))

model.add(Dense(1))
model.add(Activation('sigmoid'))

model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

img_gen = ImageDataGenerator(rescale=1./255)
train_gen = img_gen.flow_from_directory(train_folder, target_size=(img_w,img_h),
    batch_size=batch_size, class_mode='binary')
test_gen = img_gen.flow_from_directory(test_folder, target_size=(img_w,img_h),
    batch_size=batch_size, class_mode='binary')

model.fit_generator(train_gen, steps_per_epoch=train_samples//batch_size,epochs=epochs,
    validation_data=test_gen)

if not os.path.isdir(save_dir):
    os.makedirs(save_dir)

model_path = os.path.join(save_dir, 'cvd')
model.save(model_path)
print('Saved trained model at %s ' % model_path)