## Author: SHIVANK SHARMA
## Modified: Toshio Iwata
## Released under Apache 2.0 license 

import numpy as np
import pandas as pd
import os
import cv2
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
from PIL import Image
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
from sklearn.metrics import accuracy_score
np.random.seed(42)

from matplotlib import style
style.use('fivethirtyeight')

## Modified paths, be careful of the location of Train and Test folder
data_dir = './'
train_path = './Train'
test_path = './'

# Resizing the images to 30x30x3
IMG_HEIGHT = 30
IMG_WIDTH = 30
channels = 3

NUM_CATEGORIES = len(os.listdir(train_path))
NUM_CATEGORIES

# Label Overview
classes = { 0:'Speed limit (20km/h)',
            1:'Speed limit (30km/h)', 
            2:'Speed limit (50km/h)', 
            3:'Speed limit (60km/h)', 
            4:'Speed limit (70km/h)', 
            5:'Speed limit (80km/h)', 
            6:'End of speed limit (80km/h)', 
            7:'Speed limit (100km/h)', 
            8:'Speed limit (120km/h)', 
            9:'No passing' }

folders = os.listdir(train_path)

train_number = []
class_num = []

for folder in folders:
    train_files = os.listdir(train_path + '/' + folder)
    train_number.append(len(train_files))
    class_num.append(classes[int(folder)])
    
# Sorting the dataset on the basis of number of images in each class
zipped_lists = zip(train_number, class_num)
sorted_pairs = sorted(zipped_lists)
tuples = zip(*sorted_pairs)
train_number, class_num = [ list(tuple) for tuple in  tuples]

# Plotting the number of images in each class
plt.figure(figsize=(21,10))  
plt.bar(class_num, train_number)
plt.xticks(class_num, rotation='vertical')
plt.show()


image_data = []
image_labels = []

for i in range(NUM_CATEGORIES):
    path = data_dir + '/Train/' + str(i)
    images = os.listdir(path)

    for img in images:
        try:
            image = cv2.imread(path + '/' + img)
            image_fromarray = Image.fromarray(image, 'RGB')
            resize_image = image_fromarray.resize((IMG_HEIGHT, IMG_WIDTH))
            image_data.append(np.array(resize_image))
            image_labels.append(i)
        except:
            print("Error in " + img)

# Changing the list to numpy array
image_data = np.array(image_data)
image_labels = np.array(image_labels)

print(image_data.shape, image_labels.shape)

shuffle_indexes = np.arange(image_data.shape[0])
np.random.shuffle(shuffle_indexes)
image_data = image_data[shuffle_indexes]
image_labels = image_labels[shuffle_indexes]

X_train, X_val, y_train, y_val = train_test_split(image_data, image_labels, test_size=0.3, random_state=42, shuffle=True)

X_train = X_train/255 
X_val = X_val/255

print("X_train.shape", X_train.shape)
print("X_valid.shape", X_val.shape)
print("y_train.shape", y_train.shape)
print("y_valid.shape", y_val.shape)

y_train = keras.utils.to_categorical(y_train, NUM_CATEGORIES)
y_val = keras.utils.to_categorical(y_val, NUM_CATEGORIES)

print(y_train.shape)
print(y_val.shape)

model = keras.models.Sequential([    
    keras.layers.Conv2D(filters=16, kernel_size=(3,3), activation='relu', input_shape=(IMG_HEIGHT,IMG_WIDTH,channels)),
    keras.layers.Conv2D(filters=32, kernel_size=(3,3), activation='relu'),
    keras.layers.MaxPool2D(pool_size=(2, 2)),
    ##keras.layers.BatchNormalization(axis=-1),
    
    keras.layers.Conv2D(filters=64, kernel_size=(3,3), activation='relu'),
    keras.layers.Conv2D(filters=128, kernel_size=(3,3), activation='relu'),
    keras.layers.MaxPool2D(pool_size=(2, 2)),
    ##keras.layers.BatchNormalization(axis=-1),
    
    keras.layers.Flatten(),
    ##keras.layers.Dense(512, activation='relu'),
    ##keras.layers.Dense(256, activation='relu'),
    ## Modified from 512 to 128
    keras.layers.Dense(128, activation='relu'),
    ##keras.layers.BatchNormalization(),
    keras.layers.Dropout(rate=0.5),
    
    ###keras.layers.Dense(43, activation='softmax')
    keras.layers.Dense(10, activation='softmax')
])

lr = 0.001
epochs = 30

####opt = Adam(lr=lr, decay=lr / (epochs * 0.5)) for older tensorflow
opt = tf.keras.optimizers.legacy.Adam(lr=lr, decay=lr / (epochs * 0.5)) # for newer tensorflow

model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])

aug = ImageDataGenerator(
    rotation_range=10,
    zoom_range=0.15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.15,
    horizontal_flip=False,
    vertical_flip=False,
    fill_mode="nearest")

history = model.fit(aug.flow(X_train, y_train, batch_size=32), epochs=epochs, validation_data=(X_val, y_val))
############################### by iwata start

import csv

###################################################
## for write
###################################################
########################### lay1
lay1 = model.layers[0]
print("lay1 weight shape : ", lay1.get_weights()[0].shape)
print("lay1 bias shape : ", lay1.get_weights()[1].shape)

with open("w1.csv", mode='w') as f:
    for i in range(0, lay1.get_weights()[0].shape[2]):
        for r in range(0, lay1.get_weights()[0].shape[0]):
            for c in range(0, lay1.get_weights()[0].shape[1]):
                for k in range(0, lay1.get_weights()[0].shape[3]):
                    tmpval = lay1.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b1.csv", mode='w') as f:
    for c in range(0, lay1.get_weights()[1].shape[0]):
        tmpval = lay1.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay2
lay2 = model.layers[1]
print("lay2 weight shape : ", lay2.get_weights()[0].shape)
print("lay2 bias shape : ", lay2.get_weights()[1].shape)

with open("w2.csv", mode='w') as f:
    for i in range(0, lay2.get_weights()[0].shape[2]):
        for r in range(0, lay2.get_weights()[0].shape[0]):
            for c in range(0, lay2.get_weights()[0].shape[1]):
                for k in range(0, lay2.get_weights()[0].shape[3]):
                    tmpval = lay2.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b2.csv", mode='w') as f:
    for c in range(0, lay2.get_weights()[1].shape[0]):
        tmpval = lay2.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay3
lay3 = model.layers[3]
print("lay3 weight shape : ", lay3.get_weights()[0].shape)
print("lay3 bias shape : ", lay3.get_weights()[1].shape)

with open("w3.csv", mode='w') as f:
    for i in range(0, lay3.get_weights()[0].shape[2]):
        for r in range(0, lay3.get_weights()[0].shape[0]):
            for c in range(0, lay3.get_weights()[0].shape[1]):
                for k in range(0, lay3.get_weights()[0].shape[3]):
                    tmpval = lay3.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b3.csv", mode='w') as f:
    for c in range(0, lay3.get_weights()[1].shape[0]):
        tmpval = lay3.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay4
lay4 = model.layers[4]
print("lay4 weight shape : ", lay4.get_weights()[0].shape)
print("lay4 bias shape : ", lay4.get_weights()[1].shape)

with open("w4.csv", mode='w') as f:
    for i in range(0, lay4.get_weights()[0].shape[2]):
        for r in range(0, lay4.get_weights()[0].shape[0]):
            for c in range(0, lay4.get_weights()[0].shape[1]):
                for k in range(0, lay4.get_weights()[0].shape[3]):
                    tmpval = lay4.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b4.csv", mode='w') as f:
    for c in range(0, lay4.get_weights()[1].shape[0]):
        tmpval = lay4.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay5
lay5 = model.layers[7]
print("lay5 weight shape : ", lay5.get_weights()[0].shape)
print("lay5 bias shape : ", lay5.get_weights()[1].shape)

with open("w5.csv", mode='w') as f:
    for r in range(0, lay5.get_weights()[0].shape[0]):
        for c in range(0, lay5.get_weights()[0].shape[1]):
            tmpval = lay5.get_weights()[0][r][c] ### [0]に重み係数
            f.write(str(tmpval))
            f.write(', ')
        f.write('\n')

with open("b5.csv", mode='w') as f:
    for c in range(0, lay5.get_weights()[1].shape[0]):
        tmpval = lay5.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay6
lay6 = model.layers[9]
print("lay6 weight shape : ", lay6.get_weights()[0].shape)
print("lay6 bias shape : ", lay6.get_weights()[1].shape)

with open("w6.csv", mode='w') as f:
    for r in range(0, lay6.get_weights()[0].shape[0]):
        for c in range(0, lay6.get_weights()[0].shape[1]):
            tmpval = lay6.get_weights()[0][r][c] ### [0]に重み係数
            f.write(str(tmpval))
            f.write(', ')
        f.write('\n')

with open("b6.csv", mode='w') as f:
    for c in range(0, lay6.get_weights()[1].shape[0]):
        tmpval = lay6.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

############################### by iwata end

pd.DataFrame(history.history).plot(figsize=(8, 5))
plt.grid(True)
plt.gca().set_ylim(0, 1)
plt.show()

test = pd.read_csv(data_dir + '/Test0to9.csv')

labels = test["ClassId"].values
imgs = test["Path"].values

data =[]

for img in imgs:
    try:
        image = cv2.imread(data_dir + '/' +img)
        image_fromarray = Image.fromarray(image, 'RGB')
        resize_image = image_fromarray.resize((IMG_HEIGHT, IMG_WIDTH))
        data.append(np.array(resize_image))
    except:
        print("Error in " + img)
X_test = np.array(data)
X_test = X_test/255

####pred = model.predict_classes(X_test)
pred = np.argmax(model.predict(X_test), axis=-1) # for newer tensorflow

#Accuracy with the test data
print('Test Data accuracy: ',accuracy_score(labels, pred)*100)

from sklearn.metrics import confusion_matrix
cf = confusion_matrix(labels, pred)

import seaborn as sns
df_cm = pd.DataFrame(cf, index = classes,  columns = classes)
plt.figure(figsize = (20,20))
sns.heatmap(df_cm, annot=True)

from sklearn.metrics import classification_report

print(classification_report(labels, pred))

plt.figure(figsize = (25, 25))

## Modified index of images to show up, and swich the color from BGR to RGB
### changed index of images to show up
for i in range(25):
    plt.subplot(5, 5, i + 1)
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    plt.rcParams["font.size"] = 16
    if i%10 == 0 :
        rndval = int(i/10)
    elif i%10 == 1 :
        rndval = 60 + int(i/10)
    elif i%10 == 2 :
        rndval = 780 + int(i/10)
    elif i%10 == 3 :
        rndval = 1530 + int(i/10)
    elif i%10 == 4 :
        rndval = 1980 + int(i/10)
    elif i%10 == 5 :
        rndval = 2640 + int(i/10)
    elif i%10 == 6 :
        rndval = 3270 + int(i/10)
    elif i%10 == 7 :
        rndval = 3420 + int(i/10)
    elif i%10 == 8 :
        rndval = 3870 + int(i/10)
    elif i%10 == 9 :
        rndval = 4320 + int(i/10)
    prediction = pred[rndval]
    actual = labels[rndval]
    col = 'g'
    if prediction != actual:
        col = 'r'
    plt.xlabel('Actual={} || Pred={}'.format(actual, prediction), color = col)
    for m in range(0, X_test[rndval].shape[0]): 
        for n in range(0, X_test[rndval].shape[1]): 
            tmpg = X_test[rndval][m][n][0]
            tmpr = X_test[rndval][m][n][2]
            X_test[rndval][m][n][0] = tmpr
            X_test[rndval][m][n][2] = tmpg
    plt.imshow(X_test[rndval])
plt.show()