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)