def main():
# create the base pre-trained model
base_model = InceptionV3(weights='imagenet', include_top=False)
# add a global spatial average pooling layer
x = base_model.output
x = GlobalAveragePooling2D()(x)
# let's add a fully-connected layer
x = Dense(1024, activation='relu')(x)
# and a logistic layer
num_classes = len(os.listdir(INPUT_DIRECTORY))
predictions = Dense(num_classes, activation='softmax')(x)
# this is the model we will train
model = Model(inputs=base_model.input, outputs=predictions)
# first: train only the top layers (which were randomly initialized)
# i.e. freeze all convolutional InceptionV3 layers
for layer in base_model.layers:
layer.trainable = False
# compile the model (should be done *after* setting layers to non-trainable)
model.compile(optimizer='rmsprop', loss='categorical_crossentropy')
# define image generator
train_gen = image.ImageDataGenerator()
# train the model on the new data for a few epochs
model.fit_generator(train_gen.flow_from_directory(INPUT_DIRECTORY),
steps_per_epoch=3,
epochs=1,
verbose=2)
# let's predict the test set to see a rough score
labels = make_label_dict()
test_gen = image.ImageDataGenerator()
flow = test_gen.flow_from_directory(INPUT_DIRECTORY, class_mode=None)
predictions = model.predict_generator(flow, verbose=1) # steps=15611//32)
top_k = predictions.argsort()[:, -4:][:, ::-1]
classes = [" ".join([labels[i] for i in line]) for line in top_k]
filenames = flow.filenames # [os.path.basename(f) for f in flow.filenames]
csv_list = zip(filenames, classes)
write_csv(csv_list, file_name=OUTPUT_FILE)
