#************************************************************
# Image multi-classification using keras 
#************************************************************

import numpy as np
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Flatten
from keras.layers.convolutional import Conv2D
from keras.layers.convolutional import MaxPooling2D
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
import sys

np.random.seed(3)

#=============================================================================
# Data Preparation
#=============================================================================
# Data Augmentaion and use of ImageDataGenerator
train_datagen = ImageDataGenerator(rescale = 1/255.)
test1_datagen =  ImageDataGenerator(rescale = 1/255.)
test2_datagen =  ImageDataGenerator(rescale = 1/255.)

#---------------------
# Data Augmentation
#---------------------
option = ''
if len(sys.argv) == 2 :
	option = sys.argv[1]
else :
	print("\nUsage[1]: python keras-cnn-aug.py a(ugmentation)")
	print("Usage[2]: python keras-cnn-aug.py n(ormal)")
	sys.exit()

if option == 'a'  :
	print('Data Augmentation is being done......')
	data_aug_gen = ImageDataGenerator(rescale=1/255., rotation_range=10,width_shift_range=0.2,
		height_shift_range=0.2, shear_range=0.7, zoom_range=[0.9,2.2],
		horizontal_flip=True, vertical_flip=True,fill_mode='nearest')

	category = ['circle', 'rectangle', 'triangle']
	new_category = ['cir', 'rec', 'tri']
	for n in range(0,3) :					# for each folder
		for alldata_seqno in range(1,16) :	# for each original image
			#make file name
			src_folder = 'train/'+category[n]+'/'
			if alldata_seqno < 10  :
				src_img = src_folder+category[n]+'00'+str(alldata_seqno)+'.png'
			else :
				src_img = src_folder+category[n]+'0'+str(alldata_seqno)+'.png'
			print("original image : {0}...".format(src_img))

			#reshape image
			img = load_img(src_img)
			x = img_to_array(img)
			x = x.reshape((1,)+x.shape)

			# Augment 15 times per each original image
			target_folder = 'augmented/'+category[n]+'/'
			i = 0
			for batch in data_aug_gen.flow(x, batch_size=1, save_to_dir =target_folder, 
							save_prefix=new_category[n], save_format='png') :
				i += 1
				if i >= 15 :
					break

	print('\n Data Augmentation is finished')
	print('\n Copy augmented images of the augmented folder to the train folder\n')
	sys.exit()

# set data folder name
train_folder_name = 'augmented'
test1_folder_name = 'test'
test2_folder_name = 'train'
	
# make train & test generator
# https://keras.io/preprocessing/image/#flow_from_directory 참고
train_generator = train_datagen.flow_from_directory (
	train_folder_name, target_size=(24,24), batch_size=3, class_mode='categorical')
test1_generator = test1_datagen.flow_from_directory (
	test1_folder_name, target_size=(24,24), batch_size=3, class_mode='categorical',shuffle=False)
test2_generator = test2_datagen.flow_from_directory (
	test2_folder_name, target_size=(24,24), batch_size=3, class_mode='categorical',shuffle=False)

#=============================================================================
# Model Construction
#=============================================================================
model = Sequential()
model.add(Conv2D(32, kernel_size=(3,3), activation='relu',input_shape=(24,24,3)))
model.add(Conv2D(64, kernel_size=(3,3), activation='relu'))

model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Flatten())

model.add(Dense(128, activation='relu'))
model.add(Dense(3, activation='softmax'))

#=============================================================================
# Model Checking
#=============================================================================
'''
from IPython.display import SVG
from keras.utils.vis_utils import model_to_dot

%matplotlib inline

SVG(model_to_dot(model, show_shapes=True), create(prog='dot', format='svg'))
'''
#=============================================================================
# Configuring Training
#=============================================================================
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

#=============================================================================
# Model Training
#=============================================================================
'''
#In case of Small Data
model.fit_generator(train_generator, steps_per_epoch=15,epochs=50,
	validation_data=test_generator,validation_steps=5,verbose=1)
'''
#In case of Big Data
model.fit_generator(train_generator, steps_per_epoch=len(train_generator),epochs=50,
	validation_data=test1_generator,validation_steps=len(test1_generator),verbose=1)
	#steps_per_epoch=90

#=============================================================================
# Model Evaluation
#=============================================================================
print('*********** Evaluation ***********')

np.set_printoptions(formatter={'float': lambda x: '{0:0.3f}'.format(x)})

scores = model.evaluate_generator(test1_generator, steps=5)

print('%s:%.2f %%' %(model.metrics_names[1], scores[1]*100))

#=============================================================================
# Model Prediction
#=============================================================================
print('*********** Prediction ***********')

# output class mode & class indices
cls_md = "Class Mode: "+test1_generator.class_mode+"\n"
print(cls_md)
print(test1_generator.class_indices)

# output test results
print('>>>>>>>>>>  Test1 Result <<<<<<<<<<<')
output1 = model.predict_generator(test1_generator, steps=len(test1_generator), verbose=1)
print(output1)

print('>>>>>>>>>>  Test2 Result <<<<<<<<<<<')
output2 = model.predict_generator(test2_generator, steps=len(test2_generator), verbose=0)
print(output2)