classify.py

# 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)