## 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_1'
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
### Modify for floor detection
classes = { 0:'Speed limit 20km',
            1:'Speed limit 30km',
            2:'Speed limit 50km',
            3:'Speed limit 60km',
            4:'Speed limit 70km',
            5:'Speed limit 80km',
            6:'End of speed limit 80km',
            7:'Speed limit 100km',
            8:'Speed limit 120km',
            9:'No passing',
            10:'No passing veh over 3.5t',
            11:'Right-of-way at intersection',
            12:'Priority road',
            13:'Yield',
            14:'Stop',
            15:'No vehicles',
            16:'Veh over 3.5t prohibited',
            17:'No entry',
            18:'General caution',
            19:'Dangerous curve left',
            20:'Dangerous curve right',
            21:'Double curve',
            22:'Bumpy road',
            23:'Slippery road',
            24:'Road narrows on the right',
            25:'Road work',
            26:'Traffic signals',
            27:'Pedestrians',
            28:'Children crossing',
            29:'Bicycles crossing',
            30:'Beware of ice or snow',
            31:'Wild animals crossing',
            32:'End speed and passing limits',
            33:'Turn right ahead',
            34:'Turn left ahead',
            35:'Ahead only',
            36:'Go straight or right',
            37:'Go straight or left',
            38:'Keep right',
            39:'Keep left',
            40:'Roundabout mandatory',
            41:'floor' }

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()

# Visualizing 25 random images from test data
import random
from matplotlib.image import imread

'''
imgs = 'ds1_pic29.jpg' ### test["Path"].values
plt.figure(figsize=(25,25))
##plt.subplot(5,5,i)

img_path = train_path + '/0/' + imgs
rand_img = imread(img_path)
plt.imshow(rand_img)
#plt.grid(b=None)
plt.xlabel(rand_img.shape[1], fontsize = 20)#width of image


test = pd.read_csv(data_dir + '/Test.csv')
imgs = test["Path"].values

plt.figure(figsize=(25,25))
for i in range(1,26):
    plt.subplot(5,5,i)
    random_img_path = data_dir + '/' + random.choice(imgs)
    rand_img = imread(random_img_path)
    plt.imshow(rand_img)
    plt.grid(b=None)
    plt.xlabel(rand_img.shape[1], fontsize = 20)#width of image
   
'''


image_data = []
image_labels = []

for i in range(NUM_CATEGORIES):
    path = data_dir + '/train_1/' + 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(42, activation='softmax') ### Modify for floor detection
])

lr = 0.001
epochs = 10

####opt = Adam(lr=lr, decay=lr / (epochs * 0.5)) # for older jupyter lab
###opt = tf.keras.optimizers.legacy.Adam(lr=lr, decay=lr / (epochs * 0.5)) # for newer jupyter lab
opt = tf.keras.optimizers.Adam(learning_rate=lr) # for raspi 2024/07/31

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 weight 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 weight 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 weight 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 weight 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 weight 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 weight 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()