
import cv2
import numpy as np
import matplotlib.pyplot as plt

from keras.models import Sequential
from keras.preprocessing import image
from keras.preprocessing.image import ImageDataGenerator
from keras.layers import Activation, Dropout, Flatten, Conv2D, MaxPooling2D, Dense

cat4 = cv2.imread('../data/CATS_DOGS/train/CAT/4.jpg')
cat4 = cv2.cvtColor(cat4, cv2.COLOR_BGR2RGB)
# plt.imshow(cat4)
# plt.show()

image_gen = ImageDataGenerator(rotation_range=30, width_shift_range=0.1, height_shift_range=0.1, rescale=1/255,
                               shear_range=0.2, zoom_range=0.2, horizontal_flip=True, fill_mode='nearest')

# plt.imshow(image_gen.random_transform(cat4))
# plt.show()

image_gen.flow_from_directory('../data/CATS_DOGS/train')

input_shape = (150, 150, 3)

model = Sequential()
model.add(Conv2D(filters=32, kernel_size=(3, 3), input_shape=(150, 150, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(filters=64, kernel_size=(3, 3), input_shape=(150, 150, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(filters=64, kernel_size=(3, 3), input_shape=(150, 150, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Flatten())

model.add(Dense(128))
model.add(Activation('relu'))

# Help for overfiting (drop 50% of neurons)
model.add(Dropout(0.5))

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

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

print(model.summary())

batch_size = 16

training_image_gen = image_gen.flow_from_directory('../data/CATS_DOGS/train', target_size=input_shape[:2],
                                                   batch_size=batch_size, class_mode='binary')

test_image_gen = image_gen.flow_from_directory('../data/CATS_DOGS/test', target_size=input_shape[:2],
                                               batch_size=batch_size, class_mode='binary')

results = model.fit_generator(training_image_gen, epochs=1, steps_per_epoch=150, validation_data=test_image_gen,
                              validation_steps=12)

print(results.history['acc'])

dog_file = '../data/CATS_DOGS/test/DOG/10005.jpg'
dog_image = image.load_img(dog_file, target_size=(150, 150))
dog_image = image.img_to_array(dog_image)

dog_image = np.expand_dims(dog_image, axis=0)
dog_image = dog_image / 255

print(model.predict_classes(dog_image))
print(model.predict(dog_image))

