Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
# Code inspired by (or copied from)
# https://github.com/fchollet/keras/blob/master/keras/applications/vgg16.py
# https://gist.github.com/baraldilorenzo/07d7802847aaad0a35d3
#
# VGG16 weights: https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels.h5
# VGG16 image classes: https://github.com/raghakot/keras-vis/blob/master/resources/imagenet_class_index.json
#
import json
import sys
import numpy as np
from PIL import Image
from matplotlib import pyplot
from keras.models import Sequential
from keras.layers.core import Flatten, Dense, Dropout
from keras.layers.convolutional import Conv2D, MaxPooling2D, ZeroPadding2D
from keras.optimizers import SGD
from keras.preprocessing import image
import time
# Load the test image
img_path = sys.argv[1]
img = Image.open(img_path)
if img.mode == 'P':
img = img.convert('RGBA')
background = Image.new('RGBA', img.size, (255,255,255))
img = Image.alpha_composite(background, img)
if img.mode != 'RGB':
img = img.convert('RGB')
img = img.resize((224, 224))
pyplot.imshow(img)
pyplot.show()
# Create model
model = Sequential()
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', input_shape=(224,224,3)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2,2), strides=(2,2)))
model.add(Conv2D(128,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(128,( 3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2,2), strides=(2,2)))
model.add(Conv2D(256,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(256,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(256,( 3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2,2), strides=(2,2)))
model.add(Conv2D(512,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(512,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(512,( 3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2,2), strides=(2,2)))
model.add(Conv2D(512,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(512,( 3, 3), activation='relu', padding='same'))
model.add(Conv2D(512,( 3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D((2,2), strides=(2,2)))
model.add(Flatten())
model.add(Dense(4096, activation='relu'))
model.add(Dense(4096, activation='relu'))
model.add(Dense(1000, activation='softmax'))
# Load the pre-trained weights
model.load_weights('vgg16_weights_tf_dim_ordering_tf_kernels.h5')
# Feed the image into the model
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = x[:, :, :, ::-1]
x[:, :, :, 0] -= 103.939
x[:, :, :, 1] -= 116.779
x[:, :, :, 2] -= 123.68
predictions = model.predict(x)
# Load the classification index
classificationIndex = json.load(open('imagenet_class_index.json'))
# Print the top 5 classification results
results = []
for pred in predictions:
top_indices = pred.argsort()[-5:][::-1]
result = [tuple(classificationIndex[str(i)]) + (pred[i],) for i in top_indices]
result.sort(key=lambda x: x[2], reverse=True)
results.append(result)
print('results', results)
time.sleep(5)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment