├── trained_model └── challenge1.weights ├── consts.py ├── README.md ├── pred.py ├── dataHandler.py └── train.py /trained_model/challenge1.weights: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mynameisguy/TrafficLightChallenge-DeepLearning-Nexar/HEAD/trained_model/challenge1.weights -------------------------------------------------------------------------------- /consts.py: -------------------------------------------------------------------------------- 1 | ## dirs confs 2 | DATASET_FOLDER="./dataset" 3 | MODELS_CHECKPOINTS_DIR="./checkpoints" 4 | 5 | # training confs 6 | BATCH_SIZE = 64 7 | TRAINING_EPOCHS = 200 #max 8 | TRAIN_IMAGES_PER_EPOCH = 16768 9 | VALIDATE_IMAGES_PER_EPOCH = 1856 10 | IMAGE_WIDTH = 224 11 | IMAGE_HEIGHT = 224 12 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Traffic Light Recognition Network 2 | 3 | This model was created for the [nexar](https://www.getnexar.com/) challenge 1. 4 | It reached 93.6% accuracy on the challenge test data. 5 | 6 | This work is based on the [squeezeNet](https://arxiv.org/abs/1602.07360) 7 | model and it's [kera's implementation](https://github.com/DT42/squeezenet_demo). 8 | 9 | # Using 10 | in order to use the model to classify image to 3 classes: 11 | 12 | - no traffic light 13 | - red light 14 | - green light 15 | 16 | just edit the variable `TEST_FOLDER` in pred.py file. 17 | then: 18 | ``` 19 | python pred.py 20 | ``` 21 | 22 | it will create csv file: `results.csv` 23 | 24 | 25 | ## Auther 26 | [Guy Hadash](https://www.linkedin.com/in/guy-hadash-b86a8031) -------------------------------------------------------------------------------- /pred.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from keras.preprocessing.image import load_img, img_to_array 4 | import numpy as np 5 | import keras.backend as K 6 | from train import SqueezeNet 7 | from consts import IMAGE_WIDTH, IMAGE_HEIGHT 8 | 9 | TEST_FOLDER="./test" 10 | K.set_image_dim_ordering('tf') 11 | 12 | model = SqueezeNet(3, (IMAGE_HEIGHT, IMAGE_WIDTH, 3)) 13 | model.load_weights("trained_model/challenge1.weights") 14 | 15 | with open('results.csv','w') as results: 16 | results.write("fname,class,p0,p1,p2\n") 17 | for root, dirnames, filenames in os.walk(TEST_FOLDER): 18 | for filename in filenames: 19 | filepath = os.path.join(root, filename) 20 | image = load_img(filepath, target_size=(224, 224)) 21 | image = img_to_array(image) 22 | image /= 255.0 23 | image = np.expand_dims(image, axis=0) 24 | preds = model.predict(image)[0] 25 | line = filename + "," + str(np.argmax(preds)) 26 | for i in range(3): 27 | line += "," + str(preds[i]) 28 | 29 | results.write(line+'\n') 30 | 31 | -------------------------------------------------------------------------------- /dataHandler.py: -------------------------------------------------------------------------------- 1 | from keras.preprocessing import image 2 | from consts import DATASET_FOLDER, BATCH_SIZE, IMAGE_HEIGHT, IMAGE_WIDTH 3 | 4 | 5 | class dataHandler: 6 | def getGenerators(self, batch_size=BATCH_SIZE): 7 | datagen = image.ImageDataGenerator(rescale=1.0/255, shear_range=0.2, 8 | zoom_range=0.2, 9 | rotation_range=5, 10 | horizontal_flip=True) 11 | 12 | val_datagen = image.ImageDataGenerator(rescale=1.0/255) 13 | 14 | 15 | print("creating train generator") 16 | train_generator = datagen.flow_from_directory( 17 | DATASET_FOLDER + '/train', 18 | target_size=(IMAGE_HEIGHT, IMAGE_WIDTH), 19 | batch_size=batch_size, 20 | class_mode='categorical') 21 | 22 | print("creating validation generator") 23 | validation_generator = val_datagen.flow_from_directory( 24 | DATASET_FOLDER + '/test', 25 | target_size=(IMAGE_HEIGHT, IMAGE_WIDTH), 26 | batch_size=batch_size, 27 | class_mode='categorical') 28 | 29 | return train_generator, validation_generator 30 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import keras 2 | from keras.models import Model 3 | from keras.layers import Input, Activation, merge 4 | from keras.layers import Flatten, Dropout 5 | from keras.layers import Convolution2D, MaxPooling2D 6 | from keras.layers import AveragePooling2D 7 | from keras.optimizers import Adam 8 | import keras.backend as K 9 | 10 | 11 | from dataHandler import dataHandler 12 | 13 | from consts import * 14 | 15 | 16 | K.set_image_dim_ordering('tf') 17 | 18 | def SqueezeNet(nb_classes, inputs=(224, 224, 3)): 19 | """ Keras Implementation of SqueezeNet(arXiv 1602.07360) 20 | 21 | @param nb_classes: total number of final categories 22 | 23 | Arguments: 24 | inputs -- shape of the input images (channel, cols, rows) 25 | 26 | """ 27 | 28 | input_img = Input(shape=inputs) 29 | conv1 = Convolution2D( 30 | 96, 7, 7, activation='relu', init='glorot_uniform', 31 | subsample=(2, 2), border_mode='same', name='conv1')(input_img) 32 | maxpool1 = MaxPooling2D( 33 | pool_size=(3, 3), strides=(2, 2), name='maxpool1')(conv1) 34 | 35 | fire2_squeeze = Convolution2D( 36 | 16, 1, 1, activation='relu', init='glorot_uniform', 37 | border_mode='same', name='fire2_squeeze')(maxpool1) 38 | fire2_expand1 = Convolution2D( 39 | 64, 1, 1, activation='relu', init='glorot_uniform', 40 | border_mode='same', name='fire2_expand1')(fire2_squeeze) 41 | fire2_expand2 = Convolution2D( 42 | 64, 3, 3, activation='relu', init='glorot_uniform', 43 | border_mode='same', name='fire2_expand2')(fire2_squeeze) 44 | merge2 = merge( 45 | [fire2_expand1, fire2_expand2], mode='concat', concat_axis=-1) 46 | 47 | fire3_squeeze = Convolution2D( 48 | 16, 1, 1, activation='relu', init='glorot_uniform', 49 | border_mode='same', name='fire3_squeeze')(merge2) 50 | fire3_expand1 = Convolution2D( 51 | 64, 1, 1, activation='relu', init='glorot_uniform', 52 | border_mode='same', name='fire3_expand1')(fire3_squeeze) 53 | fire3_expand2 = Convolution2D( 54 | 64, 3, 3, activation='relu', init='glorot_uniform', 55 | border_mode='same', name='fire3_expand2')(fire3_squeeze) 56 | merge3 = merge( 57 | [fire3_expand1, fire3_expand2], mode='concat', concat_axis=-1) 58 | 59 | fire4_squeeze = Convolution2D( 60 | 32, 1, 1, activation='relu', init='glorot_uniform', 61 | border_mode='same', name='fire4_squeeze')(merge3) 62 | fire4_expand1 = Convolution2D( 63 | 128, 1, 1, activation='relu', init='glorot_uniform', 64 | border_mode='same', name='fire4_expand1')(fire4_squeeze) 65 | fire4_expand2 = Convolution2D( 66 | 128, 3, 3, activation='relu', init='glorot_uniform', 67 | border_mode='same', name='fire4_expand2')(fire4_squeeze) 68 | merge4 = merge( 69 | [fire4_expand1, fire4_expand2], mode='concat', concat_axis=-1) 70 | maxpool4 = MaxPooling2D( 71 | pool_size=(3, 3), strides=(2, 2), name='maxpool4')(merge4) 72 | 73 | fire5_squeeze = Convolution2D( 74 | 32, 1, 1, activation='relu', init='glorot_uniform', 75 | border_mode='same', name='fire5_squeeze')(maxpool4) 76 | fire5_expand1 = Convolution2D( 77 | 128, 1, 1, activation='relu', init='glorot_uniform', 78 | border_mode='same', name='fire5_expand1')(fire5_squeeze) 79 | fire5_expand2 = Convolution2D( 80 | 128, 3, 3, activation='relu', init='glorot_uniform', 81 | border_mode='same', name='fire5_expand2')(fire5_squeeze) 82 | merge5 = merge( 83 | [fire5_expand1, fire5_expand2], mode='concat', concat_axis=-1) 84 | 85 | fire6_squeeze = Convolution2D( 86 | 48, 1, 1, activation='relu', init='glorot_uniform', 87 | border_mode='same', name='fire6_squeeze')(merge5) 88 | fire6_expand1 = Convolution2D( 89 | 192, 1, 1, activation='relu', init='glorot_uniform', 90 | border_mode='same', name='fire6_expand1')(fire6_squeeze) 91 | fire6_expand2 = Convolution2D( 92 | 192, 3, 3, activation='relu', init='glorot_uniform', 93 | border_mode='same', name='fire6_expand2')(fire6_squeeze) 94 | merge6 = merge( 95 | [fire6_expand1, fire6_expand2], mode='concat', concat_axis=-1) 96 | 97 | fire7_squeeze = Convolution2D( 98 | 48, 1, 1, activation='relu', init='glorot_uniform', 99 | border_mode='same', name='fire7_squeeze')(merge6) 100 | fire7_expand1 = Convolution2D( 101 | 192, 1, 1, activation='relu', init='glorot_uniform', 102 | border_mode='same', name='fire7_expand1')(fire7_squeeze) 103 | fire7_expand2 = Convolution2D( 104 | 192, 3, 3, activation='relu', init='glorot_uniform', 105 | border_mode='same', name='fire7_expand2')(fire7_squeeze) 106 | merge7 = merge( 107 | [fire7_expand1, fire7_expand2], mode='concat', concat_axis=-1) 108 | 109 | fire8_squeeze = Convolution2D( 110 | 64, 1, 1, activation='relu', init='glorot_uniform', 111 | border_mode='same', name='fire8_squeeze')(merge7) 112 | fire8_expand1 = Convolution2D( 113 | 256, 1, 1, activation='relu', init='glorot_uniform', 114 | border_mode='same', name='fire8_expand1')(fire8_squeeze) 115 | fire8_expand2 = Convolution2D( 116 | 256, 3, 3, activation='relu', init='glorot_uniform', 117 | border_mode='same', name='fire8_expand2')(fire8_squeeze) 118 | merge8 = merge( 119 | [fire8_expand1, fire8_expand2], mode='concat', concat_axis=-1) 120 | 121 | maxpool8 = MaxPooling2D( 122 | pool_size=(3, 3), strides=(2, 2), name='maxpool8')(merge8) 123 | 124 | fire9_squeeze = Convolution2D( 125 | 64, 1, 1, activation='relu', init='glorot_uniform', 126 | border_mode='same', name='fire9_squeeze')(maxpool8) 127 | fire9_expand1 = Convolution2D( 128 | 256, 1, 1, activation='relu', init='glorot_uniform', 129 | border_mode='same', name='fire9_expand1')(fire9_squeeze) 130 | fire9_expand2 = Convolution2D( 131 | 256, 3, 3, activation='relu', init='glorot_uniform', 132 | border_mode='same', name='fire9_expand2')(fire9_squeeze) 133 | merge9 = merge( 134 | [fire9_expand1, fire9_expand2], mode='concat', concat_axis=-1) 135 | 136 | fire9_dropout = Dropout(0.5, name='fire9_dropout')(merge9) 137 | conv10 = Convolution2D( 138 | nb_classes, 1, 1, init='glorot_uniform', 139 | border_mode='valid', name='conv10')(fire9_dropout) 140 | # The size should match the output of conv10 141 | 142 | avgpool10 = AveragePooling2D((13, 13), name='avgpool10')(conv10) 143 | flatten = Flatten(name='flatten')(avgpool10) 144 | softmax = Activation("softmax", name='softmax')(flatten) 145 | 146 | return Model(input=input_img, output=softmax) 147 | 148 | if __name__ == '__main__': 149 | model = SqueezeNet(3, (IMAGE_HEIGHT, IMAGE_WIDTH, 3)) 150 | 151 | dHandler = dataHandler() 152 | train_generator, validation_generator = dHandler.getGenerators() 153 | adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, decay=0.01) 154 | model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=['accuracy']) 155 | 156 | tensorBoard = keras.callbacks.TensorBoard(log_dir='./logs', histogram_freq=1, write_graph=True, write_images=True) 157 | earlyStopping = keras.callbacks.EarlyStopping(monitor='val_loss', min_delta=0, patience=12, verbose=0, mode='auto') 158 | modelCheckpoints = keras.callbacks.ModelCheckpoint(MODELS_CHECKPOINTS_DIR + 159 | "/squeezeNet_" + str(IMAGE_WIDTH) + "_" + 160 | str(IMAGE_HEIGHT) + 161 | ".{epoch:02d}-{val_loss:.2f}-{val_acc:.2f}.hdf5", 162 | monitor='val_loss', verbose=0, save_best_only=False, 163 | save_weights_only=False, mode='auto') 164 | hists = model.fit_generator(train_generator, TRAIN_IMAGES_PER_EPOCH, TRAINING_EPOCHS, 165 | validation_data=validation_generator, 166 | nb_val_samples=VALIDATE_IMAGES_PER_EPOCH, 167 | callbacks=[modelCheckpoints, tensorBoard, earlyStopping]) 168 | 169 | model.save_weights("challenge1.model") 170 | --------------------------------------------------------------------------------