diff --git a/Task 3/train_model.py b/Task 3/train_model.py new file mode 100644 index 0000000..ad5707a --- /dev/null +++ b/Task 3/train_model.py @@ -0,0 +1,52 @@ +from keras.models import Sequential +from keras.layers import Convolution2D, MaxPooling2D, Flatten, Dense + +classifier = Sequential() + +classifier.add(Convolution2D(32, (3, 3), input_shape=(64, 64, 1), activation='relu')) +classifier.add(MaxPooling2D(pool_size=(2, 2))) + +classifier.add(Convolution2D(32, (3, 3), activation='relu')) +classifier.add(MaxPooling2D(pool_size=(2, 2))) + +classifier.add(Flatten()) + +classifier.add(Dense(units=128, activation='relu')) +classifier.add(Dense(units=6, activation='softmax')) + +classifier.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy']) + + +from keras.preprocessing.image import ImageDataGenerator + +train_datagen = ImageDataGenerator( + rescale=1./255, + shear_range=0.2, + zoom_range=0.2, + horizontal_flip=True) + +test_datagen = ImageDataGenerator(rescale=1./255) #epoch + +training_set = train_datagen.flow_from_directory('data/train', + target_size=(64, 64), + batch_size=5, + color_mode='grayscale', + class_mode='categorical') + +test_set = test_datagen.flow_from_directory('data/test', + target_size=(64, 64), + batch_size=5, + color_mode='grayscale', + class_mode='categorical') + +classifier.fit_generator( + training_set, + epochs=10, + validation_data=test_set) + +#Saving +model_json = classifier.to_json() +with open("model-bw.json", "w") as json_file: + json_file.write(model_json) +classifier.save_weights('model-bw.h5') +