from keras.models import Sequential
from keras.layers.core import Flatten, Dense, Dropout
from keras.layers.convolutional import Convolution2D, MaxPooling2D, ZeroPadding2D
from keras.optimizers import SGD, Adam
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import ModelCheckpoint

import cv2, numpy as np
import pickle
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
import keras.backend as K
K.set_image_dim_ordering('th')
import theano

def pop_layer(model):
    if not model.outputs:
        raise Exception('Sequential model cannot be popped: model is empty.')

    model.layers.pop()
    if not model.layers:
        model.outputs = []
        model.inbound_nodes = []
        model.outbound_nodes = []
    else:
        model.layers[-1].outbound_nodes = []
        model.outputs = [model.layers[-1].output]
    model.built = False

batch_size = 50
train_datagen = ImageDataGenerator()
test_datagen = ImageDataGenerator()
train_generator = train_datagen.flow_from_directory('/home/kmohan/RBE595_Deep_Learning/Gesture_Recognition/Consolidated_Data_Set/train',target_size = (224,224),batch_size = batch_size)
test_generator = test_datagen.flow_from_directory('/home/kmohan/RBE595_Deep_Learning/Gesture_Recognition/Consolidated_Data_Set/test',target_size = (224,224),batch_size = batch_size)
    

model = Sequential()
model.add(ZeroPadding2D((1,1),input_shape=(3,224,224)))
conv1 = Convolution2D(64, 3, 3, activation='relu')
model.add(conv1)
model.add(ZeroPadding2D((1,1)))
conv2 = Convolution2D(64, 3, 3, activation='relu')
model.add(conv2)
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(ZeroPadding2D((1,1)))
conv3 = Convolution2D(128, 3, 3, activation='relu')
model.add(conv3)
model.add(ZeroPadding2D((1,1)))
conv4 = Convolution2D(128, 3, 3, activation='relu')
model.add(conv4)
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(ZeroPadding2D((1,1)))
conv5 = Convolution2D(256, 3, 3, activation='relu')
model.add(conv5)
model.add(ZeroPadding2D((1,1)))
conv6 = Convolution2D(256, 3, 3, activation='relu')
model.add(conv6)
model.add(ZeroPadding2D((1,1)))
conv7 = Convolution2D(256, 3, 3, activation='relu')
model.add(conv7)
model.add(ZeroPadding2D((1,1)))
conv8 = Convolution2D(256, 3, 3, activation='relu')
model.add(conv8)
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(ZeroPadding2D((1,1)))
conv9 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv9)
model.add(ZeroPadding2D((1,1)))
conv10 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv10)
model.add(ZeroPadding2D((1,1)))
conv11 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv11)
model.add(ZeroPadding2D((1,1)))
conv12 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv12)
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(ZeroPadding2D((1,1)))
conv13 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv13)
model.add(ZeroPadding2D((1,1)))
conv14 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv14)
model.add(ZeroPadding2D((1,1)))
conv15 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv15)
model.add(ZeroPadding2D((1,1)))
conv16 = Convolution2D(512, 3, 3, activation='relu')
model.add(conv16)
model.add(MaxPooling2D((2,2), strides=(2,2)))

model.add(Flatten())
model.add(Dense(4096, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(4096, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(1000,activation='relu'))
#model.layers.pop()
model.load_weights('vgg19_weights.h5')

for layer in model.layers:
	layer.trainable=False

pop_layer(model)
pop_layer(model)
model.add(Dense(5,activation='softmax'))

model.load_weights('weights_rot_trans.hdf5')
#sgd = SGD(lr=0.1, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(optimizer=Adam(lr=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])

#checkpoint = ModelCheckpoint('weights.hdf5', monitor='val_acc', verbose=1, save_best_only=True, mode='max')

#history = model.fit_generator(train_generator,samples_per_epoch = 5000,nb_epoch = 25,validation_data #= test_generator,nb_val_samples = 1000,verbose = 1 , callbacks=[checkpoint])
#print ('accuracy is',history.history['acc'])

if __name__ == "__main__":
    the_filename='gesture_label.txt'
    #with open(the_filename, 'wb') as f:
    #    pickle.dump(my_list, f)
    with open(the_filename, 'rb') as f:
        my_list = pickle.load(f)

    
    
    ########################
    # REAL-TIME PREDICTION #
    ########################

    print '... Initializing RGB stream'
    
     #### Initialize built-in webcam
    cap = cv2.VideoCapture(0)
    # Enforce size of frames
    cap.set(3, 320) 
    cap.set(4, 240)

    shot_id = 0
 
    #### Start video stream and online prediction
    while (True):
         # Capture frame-by-frame
    
#        start_time = time.clock()
        
        ret, frame = cap.read()
        
        #color_frame = color_stream.read_frame() ## VideoFrame object
        #color_frame_data = frame.get_buffer_as_uint8() ## Image buffer
        #frame = convert_frame(color_frame_data, np.uint8) ## Generate BGR frame
                
        im = cv2.resize(frame, (224, 224)).astype(np.float32)
        im[:,:,0] -= 103.939
        im[:,:,1] -= 116.779
        im[:,:,2] -= 123.68
        im = im.transpose((2,0,1))
        im = np.expand_dims(im, axis=0)
        
        out = model.predict(im)
        #print np.argmax(out)
        #print my_list[np.argmax(out)]
        
        # we need to keep in mind aspect ratio so the image does
        # not look skewed or distorted -- therefore, we calculate
        # the ratio of the new image to the old image
        #r = 100.0 / frame.shape[1]
        dim = (640, 480)
 
        # perform the actual resizing of the image and show it
        resized = cv2.resize(frame, dim, interpolation = cv2.INTER_AREA)
        
        # Write some Text
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(resized,my_list[np.argmax(out)],(20,450), font, 1, (255,255,255),1,1)
        # Display the resulting frame
        cv2.imshow('DeepNN-ABB',resized)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()
