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

📅  最后修改于: 2020-12-11 04:48:12             🧑  作者: Mango


在本课程中,您将学习使用预先训练的模型来检测给定图像中的对象。您将使用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 。因此,我们将在分类器中使用图像之前将所有图像转换为227×227尺寸。

我们还声明了一个名为mean的变量,其值为128 ,以后将用于改进分类结果。

接下来,我们将开发两个功能来处理图像。

图像处理

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

图像调整大小

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

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格式,可以输入我们的网络了。接下来,我们将加载预先训练的模型文件,并将上面的图像输入到其中以进行预测。

预测处理图像中的对象

我们首先为init设置路径,并预测在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

如您所见,该模型预测的索引值为984的对象的置信度为89% 。 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)架构并将其训练在数据集上。现在,我们将学习如何创建简单的单层NN。