In [1]:
import os
import keras
import tensorflow as tf
from keras.applications.vgg19 import VGG19
from keras.applications.vgg19 import preprocess_input, decode_predictions
from PIL import Image
import numpy as np
import os.path
import matplotlib.pyplot as plt
In [2]:
%matplotlib inline
In [3]:
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
set_session(tf.Session(config=config))
In [4]:
model = VGG19(weights='imagenet')
In [16]:
def predict_image(img):
ori_w,ori_h = img.size
new_w = 224.0;
new_h = 224.0;
if ori_w > ori_h:
bs = 224.0 / ori_h;
new_w = ori_w * bs
weight = int(new_w)
height = int(new_h)
img = img.resize( (weight, height), Image.BILINEAR )
region = ( weight / 2 - 112, 0, weight / 2 + 112, height)
img = img.crop( region )
else:
bs = 224.0 / ori_w;
new_h = ori_h * bs
weight = int(new_w)
height = int(new_h)
img = img.resize( (weight, height), Image.BILINEAR )
region = ( 0, height / 2 - 112 , weight, height / 2 + 112 )
img = img.crop( region )
x = np.array( img, dtype = 'float32' )
x[:, :, 0] = x[:, :, 0] - 123.680
x[:, :, 1] = x[:, :, 1] - 116.779
x[:, :, 2] = x[:, :, 2] - 103.939
x = np.expand_dims(x, axis=0)
results = model.predict(x)
return decode_predictions(results, top=5)[0]
In [10]:
images_path = 'test_pic'
images = os.listdir(images_path)
In [11]:
print(images)
In [18]:
for image in images:
image_path = os.path.join(images_path,image)
print(image_path)
img = plt.imread(image_path)
label = predict_image(img)
plt.title(label)
plt.imshow(img)
In [ ]: