class MultiClassDecisionDriverLayer(BaseDriverLayer):
def build_model(self):
inp = Input(shape=(200, 200, 3))
self.model = Conv2D(32, (3, 3), padding='same')(inp)
self.model = Activation('relu')(self.model)
self.model = Conv2D(32, (3, 3))(self.model)
self.model = Activation('relu')(self.model)
self.model = MaxPooling2D(pool_size=(2, 2))(self.model)
self.model = Dropout(0.25)(self.model)
self.model = Conv2D(64, (3, 3), padding='same')(self.model)
self.model = Activation('relu')(self.model)
self.model = Conv2D(64, (3, 3))(self.model)
self.model = Activation('relu')(self.model)
self.model = MaxPooling2D(pool_size=(2, 2))(self.model)
self.model = Dropout(0.25)(self.model)
self.model = Flatten()(self.model)
self.model = Dense(512)(self.model)
self.model = Activation('relu')(self.model)
self.model = Dropout(0.5)(self.model)
output1 = Dense(1, activation='sigmoid')(self.model)
output2 = Dense(1, activation='sigmoid')(self.model)
output3 = Dense(1, activation='sigmoid')(self.model)
model = Model(inp, [output1, output2, output3])
model.compile(optimizers.rmsprop(lr=0.0001, decay=1e-6),
loss=[
"binary_crossentropy",
"binary_crossentropy",
"binary_crossentropy",
],
metrics=["accuracy"])
return model
def train(self):
train_datagen = ImageDataGenerator(rescale=1. / 255,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True)
valid_datagen = ImageDataGenerator(rescale=1. / 255,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True)
test_datagen = ImageDataGenerator(rescale=1. / 255,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True)
train_generator = train_datagen.flow_from_directory(
directory=os.path.join(self.data_path, 'train/'),
target_size=(200, 200),
color_mode="rgb",
batch_size=32,
class_mode='categorical',
shuffle=True,
seed=42)
valid_generator = valid_datagen.flow_from_directory(
directory=os.path.join(self.data_path, 'valid/'),
target_size=(200, 200),
color_mode="rgb",
batch_size=32,
class_mode="categorical",
shuffle=True,
seed=42)
test_generator = test_datagen.flow_from_directory(
directory=os.path.join(self.data_path, 'test/'),
target_size=(200, 200),
color_mode="rgb",
batch_size=32,
class_mode='categorical')
STEP_SIZE_TRAIN = train_generator.n // train_generator.batch_size
STEP_SIZE_VALID = valid_generator.n // valid_generator.batch_size
self.model.fit_generator(generator=train_generator,
steps_per_epoch=STEP_SIZE_TRAIN,
validation_data=valid_generator,
validation_steps=STEP_SIZE_VALID,
epochs=10)
self.model.evaluate_generator(generator=valid_generator,
steps=STEP_SIZE_VALID)
STEP_SIZE_TEST = test_generator.n // test_generator.batch_size
test_generator.reset()
pred = self.model.predict_generator(test_generator,
steps=STEP_SIZE_TEST,
verbose=1)
predicted_class_indices = np.argmax(pred, axis=1)
labels = train_generator.class_indices
self.labels = dict((v, k) for k, v in labels.items())
predictions = [self.labels[k] for k in predicted_class_indices]
print(predictions)
def predict(self, image_data):
image_array = image.img_to_array(image_data)
image_array = np.expand_dims(image_array, axis=0)
result = self.model.predict_classes(image_array)
return self.labels.get(result[0])
