使用预训练模型进行图像分类


在本课程中,您将学习使用预训练模型来检测给定图像中的对象。您将使用squeezenet预训练模块来高精度地检测和分类给定图像中的对象。

打开新的Juypter 笔记本并按照步骤开发此图像分类应用程序。

导入库

首先,我们使用以下代码导入所需的包 -

from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace, models
import numpy as np
import skimage.io
import skimage.transform
from matplotlib import pyplot
import os
import urllib.request as urllib2
import operator

接下来,我们设置一些变量-

INPUT_IMAGE_SIZE = 227
mean = 128

用于训练的图像显然具有不同的尺寸。所有这些图像必须转换为固定尺寸才能进行准确的训练。同样,测试图像和生产环境中要预测的图像也必须转换为与训练时使用的尺寸相同的尺寸。因此,我们在上面创建了一个名为INPUT_IMAGE_SIZE 的变量,其值为227。因此,在我们的分类器中使用之前,我们会将所有图像转换为227x227 的大小。

我们还声明了一个名为Mean 的变量,其值为128,稍后将使用该变量来改进分类结果。

接下来,我们将开发两个用于处理图像的函数。

图像处理

图像处理由两个步骤组成。第一个是调整图像大小,第二个是集中裁剪图像。对于这两个步骤,我们将编写两个用于调整大小和裁剪的函数。

调整图像大小

首先,我们将编写一个用于调整图像大小的函数。如前所述,我们将图像大小调整为227x227。因此,让我们定义函数resize如下 -

def resize(img, input_height, input_width):

我们通过将宽度除以高度来获得图像的纵横比。

original_aspect = img.shape[1]/float(img.shape[0])

如果长宽比大于1,则表明图像很宽,即处于横向模式。我们现在调整图像高度并使用以下代码返回调整大小的图像 -

if(original_aspect>1):
   new_height = int(original_aspect * input_height)
   return skimage.transform.resize(img, (input_width,
   new_height), mode='constant', anti_aliasing=True, anti_aliasing_sigma=None)

如果长宽比小于1,则表示纵向模式。我们现在使用以下代码调整宽度 -

if(original_aspect<1):
   new_width = int(input_width/original_aspect)
   return skimage.transform.resize(img, (new_width,
   input_height), mode='constant', anti_aliasing=True, anti_aliasing_sigma=None)

如果长宽比等于1,我们不进行任何高度/宽度调整。

if(original_aspect == 1):
   return skimage.transform.resize(img, (input_width,
   input_height), mode='constant', anti_aliasing=True, anti_aliasing_sigma=None)

下面给出了完整的功能代码供您快速参考 -

def resize(img, input_height, input_width):
   original_aspect = img.shape[1]/float(img.shape[0])
   if(original_aspect>1):
      new_height = int(original_aspect * input_height)
      return skimage.transform.resize(img, (input_width,
	   new_height), mode='constant', anti_aliasing=True, anti_aliasing_sigma=None)
   if(original_aspect<1):
         new_width = int(input_width/original_aspect)
         return skimage.transform.resize(img, (new_width,
         input_height), mode='constant', anti_aliasing=True, anti_aliasing_sigma=None)
   if(original_aspect == 1):
         return skimage.transform.resize(img, (input_width,
         input_height), mode='constant', anti_aliasing=True, anti_aliasing_sigma=None)

我们现在将编写一个函数来围绕图像中心裁剪图像。

图像裁剪

我们声明crop_image函数如下:

def crop_image(img,cropx,cropy):

我们使用以下语句提取图像的尺寸 -

y,x,c = img.shape

我们使用以下两行代码为图像创建一个新的起点 -

startx = x//2-(cropx//2)
starty = y//2-(cropy//2)

最后,我们通过创建具有新尺寸的图像对象来返回裁剪后的图像 -

return img[starty:starty+cropy,startx:startx+cropx]

下面给出了完整的函数代码供您快速参考 -

def crop_image(img,cropx,cropy):
   y,x,c = img.shape
   startx = x//2-(cropx//2)
   starty = y//2-(cropy//2)
   return img[starty:starty+cropy,startx:startx+cropx]

现在,我们将编写代码来测试这些功能。

处理图像

首先,将图像文件复制到项目目录中的images 子文件夹中。tree.jpg文件被复制到项目中。以下 Python 代码加载图像并将其显示在控制台上 -

img = skimage.img_as_float(skimage.io.imread("images/tree.jpg")).astype(np.float32)
print("Original Image Shape: " , img.shape)
pyplot.figure()
pyplot.imshow(img)
pyplot.title('Original image')

输出如下 -

处理图像

请注意,原始图像的尺寸为600 x 960。我们需要将其大小调整为我们的规格227 x 227。调用我们之前定义的调整大小函数可以完成这项工作。

img = resize(img, INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE)
print("Image Shape after resizing: " , img.shape)
pyplot.figure()
pyplot.imshow(img)
pyplot.title('Resized image')

输出如下 -

原始图像

请注意,现在图像大小为227 x 363。我们需要将其裁剪为227 x 227,作为我们算法的最终输入。为此,我们调用之前定义的裁剪函数。

img = crop_image(img, INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE)
print("Image Shape after cropping: " , img.shape)
pyplot.figure()
pyplot.imshow(img)
pyplot.title('Center Cropped')

下面提到的是代码的输出 -

裁剪图像

此时,图像的大小为227 x 227,可以进行进一步处理。现在,我们交换图像轴以将三种颜色提取到三个不同的区域。

img = img.swapaxes(1, 2).swapaxes(0, 1)
print("CHW Image Shape: " , img.shape)

下面给出的是输出 -

CHW Image Shape: (3, 227, 227)

请注意,最后一个轴现在已成为数组中的第一个维度。我们现在将使用以下代码绘制三个通道 -

pyplot.figure()
for i in range(3):
   pyplot.subplot(1, 3, i+1)
   pyplot.imshow(img[i])
   pyplot.axis('off')
   pyplot.title('RGB channel %d' % (i+1))

输出如下:

方面

最后,我们对图像进行一些额外的处理,例如将红绿蓝转换为蓝绿红(RGB 到 BGR),删除平均值以获得更好的结果,并使用以下三行代码添加批量大小轴 -

# convert RGB --> BGR
img = img[(2, 1, 0), :, :]
# remove mean
img = img * 255 - mean
# add batch size axis
img = img[np.newaxis, :, :, :].astype(np.float32)

此时,您的图像已采用NCHW 格式,并已准备好输入我们的网络。接下来,我们将加载预先训练的模型文件并将上面的图像输入其中进行预测。

预测处理图像中的对象

我们首先设置Caffe 预训练模型中定义的初始化预测网络的路径。

设置模型文件路径

请记住,从我们之前的讨论来看,所有预训练的模型都安装在models文件夹中。我们设置该文件夹的路径如下 -

CAFFE_MODELS = os.path.expanduser("/anaconda3/lib/python3.7/site-packages/caffe2/python/models")

我们设置squeezenet模型的init_net protobuf文件的路径如下 -

INIT_NET = os.path.join(CAFFE_MODELS, 'squeezenet', 'init_net.pb')

同样,我们设置Predict_net protobuf 的路径如下 -

PREDICT_NET = os.path.join(CAFFE_MODELS, 'squeezenet', 'predict_net.pb')

我们打印两条路径用于诊断目的 -

print(INIT_NET)
print(PREDICT_NET)

此处给出上述代码以及输出以供您快速参考 -

CAFFE_MODELS = os.path.expanduser("/anaconda3/lib/python3.7/site-packages/caffe2/python/models")
INIT_NET = os.path.join(CAFFE_MODELS, 'squeezenet', 'init_net.pb')
PREDICT_NET = os.path.join(CAFFE_MODELS, 'squeezenet', 'predict_net.pb')
print(INIT_NET)
print(PREDICT_NET)

输出如下 -

/anaconda3/lib/python3.7/site-packages/caffe2/python/models/squeezenet/init_net.pb
/anaconda3/lib/python3.7/site-packages/caffe2/python/models/squeezenet/predict_net.pb

接下来,我们将创建一个预测器。

创建预测器

我们使用以下两个语句读取模型文件 -

with open(INIT_NET, "rb") as f:
   init_net = f.read()
with open(PREDICT_NET, "rb") as f:
   predict_net = f.read()

通过将两个文件的指针作为参数传递给Predictor函数来创建预测器。

p = workspace.Predictor(init_net, predict_net)

p对象是预测器,用于预测任何给定图像中的对象。请注意,每个输入图像必须采用 NCHW 格式,就像我们之前对tree.jpg文件所做的那样。

预测物体

预测给定图像中的对象很简单 - 只需执行一行命令即可。我们在预测器对象上调用run方法来检测给定图像中的对象。

results = p.run({'data': img})

预测结果现在在结果对象中可用,为了便于阅读,我们将其转换为数组。

results = np.asarray(results)

使用以下语句打印数组的维度以供您理解 -

print("results shape: ", results.shape)

输出如下所示 -

results shape: (1, 1, 1000, 1, 1)

我们现在将删除不必要的轴 -

preds = np.squeeze(results)

现在可以通过获取preds数组中的最大值检索最上面的预测。

curr_pred, curr_conf = max(enumerate(preds), key=operator.itemgetter(1))
print("Prediction: ", curr_pred)
print("Confidence: ", curr_conf)

输出如下 -

Prediction: 984
Confidence: 0.89235985

正如您所看到的,模型以89% 的置信度预测了一个索引值为984的对象。984的索引对于我们理解检测到什么样的物体没有多大意义。我们需要使用对象的索引值获取对象的字符串化名称。模型识别的对象类型及其相应的索引值可在 github 存储库中找到。

现在,我们将了解如何检索索引值为 984 的对象的名称。

字符串化结果

我们创建 github 存储库的 URL 对象,如下所示 -

codes = "https://gist.githubusercontent.com/aaronmarkham/cd3a6b6ac0
71eca6f7b4a6e40e6038aa/raw/9edb4038a37da6b5a44c3b5bc52e448ff09bfe5b/alexnet_codes"

我们读取 URL 的内容 -

response = urllib2.urlopen(codes)

响应将包含所有代码及其描述的列表。下面显示了几行响应,以便您理解它包含的内容 -

5: 'electric ray, crampfish, numbfish, torpedo',
6: 'stingray',
7: 'cock',
8: 'hen',
9: 'ostrich, Struthio camelus',
10: 'brambling, Fringilla montifringilla',

现在,我们使用for循环迭代整个数组以找到所需的代码 984,如下所示 -

for line in response:
   mystring = line.decode('ascii')
   code, result = mystring.partition(":")[::2]
   code = code.strip()
   result = result.replace("'", "")
   if (code == str(curr_pred)):
      name = result.split(",")[0][1:]
      print("Model predicts", name, "with", curr_conf, "confidence")

当您运行代码时,您将看到以下输出 -

Model predicts rapeseed with 0.89235985 confidence

您现在可以在另一张图像上尝试该模型。

预测不同的图像

要预测另一个图像,只需将图像文件复制到项目目录的images文件夹中即可。这是我们之前存储的tree.jpg文件的目录。更改代码中图像文件的名称。只需要进行一处更改,如下所示

img = skimage.img_as_float(skimage.io.imread("images/pretzel.jpg")).astype(np.float32)

原始图片和预测结果如下所示 -

预测图像

输出如下 -

Model predicts pretzel with 0.99999976 confidence

正如您所看到的,预训练的模型能够非常准确地检测给定图像中的对象。

完整源码

此处提到了上述代码的完整源代码,该代码使用预先训练的模型在给定图像中进行对象检测,以供您快速参考 -

def crop_image(img,cropx,cropy):
   y,x,c = img.shape
   startx = x//2-(cropx//2)
   starty = y//2-(cropy//2)
   return img[starty:starty+cropy,startx:startx+cropx]
img = skimage.img_as_float(skimage.io.imread("images/pretzel.jpg")).astype(np.float32)
print("Original Image Shape: " , img.shape)
pyplot.figure()
pyplot.imshow(img)
pyplot.title('Original image')
img = resize(img, INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE)
print("Image Shape after resizing: " , img.shape)
pyplot.figure()
pyplot.imshow(img)
pyplot.title('Resized image')
img = crop_image(img, INPUT_IMAGE_SIZE, INPUT_IMAGE_SIZE)
print("Image Shape after cropping: " , img.shape)
pyplot.figure()
pyplot.imshow(img)
pyplot.title('Center Cropped')
img = img.swapaxes(1, 2).swapaxes(0, 1)
print("CHW Image Shape: " , img.shape)
pyplot.figure()
for i in range(3):
pyplot.subplot(1, 3, i+1)
pyplot.imshow(img[i])
pyplot.axis('off')
pyplot.title('RGB channel %d' % (i+1))
# convert RGB --> BGR
img = img[(2, 1, 0), :, :]
# remove mean
img = img * 255 - mean
# add batch size axis
img = img[np.newaxis, :, :, :].astype(np.float32)
CAFFE_MODELS = os.path.expanduser("/anaconda3/lib/python3.7/site-packages/caffe2/python/models")
INIT_NET = os.path.join(CAFFE_MODELS, 'squeezenet', 'init_net.pb')
PREDICT_NET = os.path.join(CAFFE_MODELS, 'squeezenet', 'predict_net.pb')
print(INIT_NET)
print(PREDICT_NET)
with open(INIT_NET, "rb") as f:
   init_net = f.read()
with open(PREDICT_NET, "rb") as f:
   predict_net = f.read()
p = workspace.Predictor(init_net, predict_net)
results = p.run({'data': img})
results = np.asarray(results)
print("results shape: ", results.shape)
preds = np.squeeze(results)
curr_pred, curr_conf = max(enumerate(preds), key=operator.itemgetter(1))
print("Prediction: ", curr_pred)
print("Confidence: ", curr_conf)
codes = "https://gist.githubusercontent.com/aaronmarkham/cd3a6b6ac071eca6f7b4a6e40e6038aa/raw/9edb4038a37da6b5a44c3b5bc52e448ff09bfe5b/alexnet_codes"
response = urllib2.urlopen(codes)
for line in response:
   mystring = line.decode('ascii')
   code, result = mystring.partition(":")[::2]
   code = code.strip()
   result = result.replace("'", "")
   if (code == str(curr_pred)):
      name = result.split(",")[0][1:]
      print("Model predicts", name, "with", curr_conf, "confidence")

此时,您已经知道如何使用预训练模型对数据集进行预测。

接下来是学习如何在Caffe2中定义神经网络 (NN)架构并在数据集上训练它们。我们现在将学习如何创建一个简单的单层神经网络。