## 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
from google.colab.patches import cv2_imshow

np.random.seed(42)

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

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' }

###################################################
## for making models
###################################################
# モデルを作成
model1 = keras.Sequential(name="my_model1")
# model1を確認
print(model1)
model1.add(keras.layers.Conv2D(filters=16, kernel_size=(3,3), activation='relu', input_shape=(IMG_HEIGHT,IMG_WIDTH,channels)))


model2 = keras.Sequential(name="my_model2")
# model2を確認
print(model2)
model2.add(keras.layers.Conv2D(filters=32, kernel_size=(3,3), activation='relu', input_shape=(28,28,16)))
model2.add(keras.layers.MaxPool2D(pool_size=(2, 2)))


model3 = keras.Sequential(name="my_model3")
# model3を確認
print(model3)
model3.add(keras.layers.Conv2D(filters=64, kernel_size=(3,3), activation='relu', input_shape=(13,13,32)))


model4 = keras.Sequential(name="my_model4")
# model4を確認
print(model4)
model4.add(keras.layers.Conv2D(filters=128, kernel_size=(3,3), activation='relu', input_shape=(11,11,64)))
model4.add(keras.layers.MaxPool2D(pool_size=(2, 2)))
model4.add(keras.layers.Flatten())


model5 = keras.Sequential(name="my_model5")
# model5を確認
print(model5)
##model5.add(keras.layers.Dense(256, activation='relu', input_dim=2048))
model5.add(keras.layers.Dense(128, activation='relu', input_dim=2048))
model5.add(keras.layers.Dropout(rate=0.5))


model6 = keras.Sequential(name="my_model6")
# model6を確認
print(model6)
### Modify for floor detection
model6.add(keras.layers.Dense(42, activation='softmax', input_dim=128))

import csv

###################################################
## for read
###################################################
########################### lay1
lay1 = model1.layers[0]
tmpw = lay1.get_weights()

with open('w1.csv', mode='r') as f:
    reader = csv.reader(f)
    r = 0
    c = 0
    i = 0
    k = 0
    print("w1 shape : ", lay1.get_weights()[0].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay1.get_weights()[0].shape[3]):
            tmpval = row[k] ### [0]縺ｫ驥阪∩菫よ焚
            tmpw[0][r][c][i][k] = tmpval ###row[k] ### [0]縺ｫ驥阪∩菫よ焚
        c += 1
        if(c == range(lay1.get_weights()[0].shape[1]).stop):
            c = 0
            r += 1
            if(r == range(lay1.get_weights()[0].shape[0]).stop):
                r = 0
                i += 1

with open('b1.csv', mode='r') as f:
    reader = csv.reader(f)
    print("b1 shape : ", lay1.get_weights()[1].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay1.get_weights()[1].shape[0]):
            tmpval = row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ
            tmpw[1][k] = tmpval ###row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ

lay1.set_weights(tmpw)
print("w1 read done")

########################### lay2
lay2 = model2.layers[1-1]
tmpw = lay2.get_weights()

with open('w2.csv', mode='r') as f:
    reader = csv.reader(f)
    r = 0
    c = 0
    i = 0
    k = 0
    print("w2 shape : ", lay2.get_weights()[0].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay2.get_weights()[0].shape[3]):
            tmpval = row[k] ### [0]縺ｫ驥阪∩菫よ焚
            tmpw[0][r][c][i][k] = tmpval ###row[k] ### [0]縺ｫ驥阪∩菫よ焚
        c += 1
        if(c == range(lay2.get_weights()[0].shape[1]).stop):
            c = 0
            r += 1
            if(r == range(lay2.get_weights()[0].shape[0]).stop):
                r = 0
                i += 1

with open('b2.csv', mode='r') as f:
    reader = csv.reader(f)
    print("b2 shape : ", lay2.get_weights()[1].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay2.get_weights()[1].shape[0]):
            tmpval = row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ
            tmpw[1][k] = tmpval ###row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ

lay2.set_weights(tmpw)
print("w2 read done")

########################### lay3
lay3 = model3.layers[3-3]
tmpw = lay3.get_weights()

with open('w3.csv', mode='r') as f:
    reader = csv.reader(f)
    r = 0
    c = 0
    i = 0
    k = 0
    print("w3 shape : ", lay3.get_weights()[0].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay3.get_weights()[0].shape[3]):
            tmpval = row[k] ### [0]縺ｫ驥阪∩菫よ焚
            tmpw[0][r][c][i][k] = tmpval ## row[k] ### [0]縺ｫ驥阪∩菫よ焚
        c += 1
        if(c == range(lay3.get_weights()[0].shape[1]).stop):
            c = 0
            r += 1
            if(r == range(lay3.get_weights()[0].shape[0]).stop):
                r = 0
                i += 1

with open('b3.csv', mode='r') as f:
    reader = csv.reader(f)
    print("b3 shape : ", lay3.get_weights()[1].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay3.get_weights()[1].shape[0]):
            tmpval = row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ
            tmpw[1][k] = tmpval ## row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ

lay3.set_weights(tmpw)
print("w3 read done")

########################### lay4
lay4 = model4.layers[4-4]
tmpw = lay4.get_weights()

with open('w4.csv', mode='r') as f:
    reader = csv.reader(f)
    r = 0
    c = 0
    i = 0
    k = 0
    print("w4 shape : ", lay4.get_weights()[0].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay4.get_weights()[0].shape[3]):
            tmpval = row[k] ### [0]縺ｫ驥阪∩菫よ焚
            tmpw[0][r][c][i][k] = tmpval ## row[k] ### [0]縺ｫ驥阪∩菫よ焚
        c += 1
        if(c == range(lay4.get_weights()[0].shape[1]).stop):
            c = 0
            r += 1
            if(r == range(lay4.get_weights()[0].shape[0]).stop):
                r = 0
                i += 1

with open('b4.csv', mode='r') as f:
    reader = csv.reader(f)
    print("b4 shape : ", lay4.get_weights()[1].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay4.get_weights()[1].shape[0]):
            tmpval = row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ
            tmpw[1][k] = tmpval ## row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ

lay4.set_weights(tmpw)
print("w4 read done")

########################### lay5
lay5 = model5.layers[7-7]
tmpw = lay5.get_weights()

with open('w5.csv', mode='r') as f:
    reader = csv.reader(f)
    i = 0
    k = 0
    print("w5 shape : ", lay5.get_weights()[0].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay5.get_weights()[0].shape[1]):
            tmpval = row[k] ### [0]縺ｫ驥阪∩菫よ焚
            tmpw[0][i][k] = tmpval ### row[k] ### [0]縺ｫ驥阪∩菫よ焚
        i += 1

with open('b5.csv', mode='r') as f:
    reader = csv.reader(f)
    print("b5 shape : ", lay5.get_weights()[1].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay5.get_weights()[1].shape[0]):
            tmpval = row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ
            tmpw[1][k] = tmpval ## row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ

lay5.set_weights(tmpw)
print("w5 read done")

########################### lay6
lay6 = model6.layers[9-9]
tmpw = lay6.get_weights()

with open('w6.csv', mode='r') as f:
    reader = csv.reader(f)
    i = 0
    k = 0
    print("w6 shape : ", lay6.get_weights()[0].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay6.get_weights()[0].shape[1]):
            tmpval = row[k] ### [0]縺ｫ驥阪∩菫よ焚
            tmpw[0][i][k] = tmpval ### row[k] ### [0]縺ｫ驥阪∩菫よ焚
        i += 1

with open('b6.csv', mode='r') as f:
    reader = csv.reader(f)
    print("b6 shape : ", lay6.get_weights()[1].shape)
    for row in reader:
#        print(row)
        for k in range(0, lay6.get_weights()[1].shape[0]):
            tmpval = row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ
            tmpw[1][k] = tmpval ## row[k] ### [1]縺ｫ繝舌う繧｢繧ｹ


lay6.set_weights(tmpw)
print("w6 read done")


data =[]
### Modify for floor detection
rescol =['0','1','2','3','4','5','6','7','8','9', \
         '10','11','12','13','14','15','16','17','18','19', \
         '20','21','22','23','24','25','26','27','28','29', \
         '30','31','32','33','34','35','36','37','38','39', \
         '40','F']

image = cv2.imread("withbg01.png")

kstep = 32
istep = 30
kmax = int(640/kstep)
imax = int(480/istep)

for i in range(imax):
    for k in range(kmax):
        try:
            cropped_image = image[(i+0)*istep:(i+1)*istep, (k+0)*kstep:(k+1)*kstep]
            image_fromarray = Image.fromarray(cropped_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
print("x_test shape : ", X_test.shape)

######### star prediction

lay1out = model1.predict(X_test) 
print("lay1out shape : ", lay1out.shape)
###lay1out = lay1out / 256
 
lay2out = model2.predict(lay1out) 
print("lay2out shape : ", lay2out.shape)
###lay2out = lay2out / 256

lay3out = model3.predict(lay2out) 
print("lay3out shape : ", lay3out.shape)
###lay3out = lay3out / 256

lay4out = model4.predict(lay3out) 
print("lay4out shape : ", lay4out.shape)
###lay4out = lay4out / 256

lay5out = model5.predict(lay4out) 
print("lay5out shape : ", lay5out.shape)
###lay5out = lay5out / 256

lay6out = model6.predict(lay5out) 
pred = np.argmax(lay6out, axis=-1) 

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

edframe = image
pt1 = 20, 20
pt2 = 130, 35
line_type = cv2.LINE_4
line_color = (0, 255, 0)
thickness = 1
#### kstep = 32, kmax = 20
#### istep = 30, imax = 16
for i in range(1, imax):
    pt1 = 0, i*istep
    pt2 = 640, i*istep
    cv2.line(edframe, pt1, pt2, line_color, thickness, line_type)
for k in range(1, kmax):
    pt1 = k*kstep, 0
    pt2 = k*kstep, 480
    cv2.line(edframe, pt1, pt2, line_color, thickness, line_type)

for i in range(imax):
    for k in range(kmax):
        ind = k+kmax*i
        #print('index ' + str(ind) + ': pred ' + str(pred[ind]))
        prediction = pred[ind]
        ### Modify for floor detection
        if prediction < 41:
            cv2.putText(edframe, rescol[prediction], (k*kstep+3,i*istep+23), cv2.FONT_HERSHEY_PLAIN, 1.2, (0, 255,0), 1, cv2.LINE_AA)

cv2_imshow(edframe)
###cv2.imshow('EdFrame', edframe)
cv2.waitKey(0)
cv2.destroyAllWindows()