# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Data Extraction and Preprocessing

# +
import os
import numpy as np
import pandas as pd

from __future__ import print_function 


import matplotlib.image as mpimg
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split

DATA_DIR = "./training_data"
IMAGE_DIR = "IMG"

CORRECTION = 0.2



def process_images(image_files, 
                   local_dirname, 
                   #normalizer = keras_normalizer,
                   batch_size = 512):
    '''Crops Images to remove hood and scenery above the horizon'''
    print ("Processing", len(image_files), "images")
    
    image_files = [os.path.join(DATA_DIR, local_dirname, IMAGE_DIR, os.path.basename(fn)) for fn in image_files]
    
    
    images = [mpimg.imread(f) for f in image_files]
    
    images = np.array([i[50:,][:-26] for i in images])
    
    return images 
    
    
    

def extract(dirname):
    '''extracts training images from folders generated by the simulator
       applies data augmentation by incorporating left/right camera angles with a correction factor'''
    print ("Extracting", dirname)
    log_fn = os.path.join(DATA_DIR, dirname, 'driving_log.csv')
    metadata= pd.read_csv(log_fn, 
                          sep=',', 
                          names=['center', 'left', 'right', 'angle', 'throttle', 'break', 'speed'])
    
    #print ("Metadata", metadata.head())
    
    center_images = process_images(metadata.center, dirname)
    left_images = process_images(metadata.left, dirname)
    right_images = process_images(metadata.right, dirname)
    
    # data augmentation #1 --> use multiple camera angles
    center_steering = np.array(metadata.angle)
    left_steering = np.array(metadata.angle)  + CORRECTION
    right_steering = np.array(metadata.angle) - CORRECTION
    
    
    print (center_images.shape, center_steering.shape)
    
    Xs = np.vstack((center_images,  left_images,  right_images))
    ys = np.hstack((center_steering, left_steering, right_steering))
    
    X_train, X_test, y_train, y_test = train_test_split(Xs, ys, test_size = 0.25, random_state=42)
    
    print ("Final Shapes TRAIN", X_train.shape, y_train.shape, "TEST", X_test.shape, y_test.shape)
    
    return X_train, X_test, y_train, y_test

# -

# ## VGG16-Based steering model

# +
from keras.models import Sequential, Model
from keras.layers import Flatten, Dense, Lambda, Input, Dropout
from keras import optimizers

from keras.applications import VGG16
import keras.backend.tensorflow_backend as K


def vgg16_steering_model(device='/gpu:1'):
    print ("Using Device:",device)
    #base_model = VGG16(include_top = False, weights='imagenet', input_shape=(84,320,3))
    with K.tf.device(device):
        # perform normalization at beginning of base VGG16 model
        inp = Input(shape=(84, 320, 3))
        norm = Lambda(lambda x: (x / 255.0) - 0.5 )(inp)
        base_model = VGG16(include_top = False, weights="imagenet", input_tensor=norm)

        # disable layer training 
        for layer in base_model.layers:
            #print (layer.name)
            layer.trainable = False

        # Flatten incoming data to feed to FC layers
        x = Flatten()(base_model.output)
    
        # add some FC layers with dropout to learn from the features detected by VGG 16
        x = Dense(256, activation="relu")(x)
        x = Dropout(0.1)(x)
        x = Dense(128, activation="relu")(x)
        x = Dropout(0.1)(x)
        
        prediction = Dense(1)(x)

        head_model = Model(input = base_model.input, output = prediction )

        opt = optimizers.Adam(lr=0.001)

        head_model.compile(loss='mse', optimizer= opt)

        head_model.summary()

        return head_model


# -

# ## Putting everything together and training the model

# +
data_dirs = os.listdir(DATA_DIR)

print (data_dirs)

BATCH_SIZE = 128
EPOCHS = 64

X_train = []
y_train = []
X_valid = []
y_valid = []

model = vgg16_steering_model()


def smart_append(X,Y, vstack=True):
    if len(X) == 0:
        return Y
    else:
        if vstack:
            return np.vstack((X,Y))
        else:
            return np.hstack((X,Y))

for d in data_dirs:
    X_tr, X_va, y_tr, y_va = extract(d)
    
    X_train = smart_append(X_train, X_tr)
    X_valid = smart_append(X_valid, X_va)
    
    y_train = smart_append(y_train, y_tr, vstack = False)
    y_valid = smart_append(y_valid, y_va, vstack = False)

print ("Training Data Shapes")
print (X_train.shape, len(y_valid))



model.fit(X_train, y_train, epochs = EPOCHS, batch_size = BATCH_SIZE)

res = model.evaluate(X_valid, y_valid)

print ("Evaluation Result", res)

import h5py
model.save('vgg_256_128.h5')
    
    
    
# -

# ## Debugging

import tensorflow as tf
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))


import sys
print (sys.version)



StarcoderdataJupyterScriptsDedupFiltered