边缘设备的实时 AI 人员检测:使用预训练 SSD 模型检测人员





5.00/5 (5投票s)
在本文中,我们将展示用于启动这些模型的 Python 代码,并检测图像中的人物。
在这一系列文章的上一篇文章中,我们为进一步工作选择了两个 SSD 模型,一个基于 MobileNet,另一个基于 SqueezeNet。在本文中,我们将开发一些 Python 代码,使我们能够使用这些模型检测图像中的人物。
选定的 DNN 以 Caffe 模型实现。Caffe 模型由两部分组成:模型结构(.prototxt)文件和训练模型(.caffemodel)。Caffe 模型结构以类似 JSON 的格式编写。训练模型是 CNN 核和其他训练数据的二进制序列化。在系列的第一篇文章中,我们提到将使用 Python OpenCV 库配合 Caffe。这是什么意思?我们是否应该同时安装这两个框架,OpenCV 和 Caffe?幸运的是,不用,只需要 OpenCV 库。该框架包含 DNN 模块,可以直接支持用 TensorFlow、Caffe、Torch、Darknet 等开发的网络模型。所以——我们很幸运!——OpenCV 框架允许同时处理计算机视觉算法和深度神经网络。这就是我们所需要的一切。
让我们用两个实用类开始我们的 Python 代码
import cv2
import numpy as np
import os
class CaffeModelLoader:
@staticmethod
def load(proto, model):
net = cv2.dnn.readNetFromCaffe(proto, model)
return net
class FrameProcessor:
def __init__(self, size, scale, mean):
self.size = size
self.scale = scale
self.mean = mean
def get_blob(self, frame):
img = frame
(h, w, c) = frame.shape
if w>h :
dx = int((w-h)/2)
img = frame[0:h, dx:dx+h]
resized = cv2.resize(img, (self.size, self.size), cv2.INTER_AREA)
blob = cv2.dnn.blobFromImage(resized, self.scale, (self.size, self.size), self.mean)
return blob
CaffeModelLoader
类有一个静态方法,用于从磁盘加载 Caffe 模型。后一个 FrameProcessor
类旨在将数据从图像转换为 DNN 指定的特定格式。类的构造函数接收三个参数。size
参数定义了 DNN 处理的输入数据的大小。用于图像处理的卷积网络几乎总是使用方形图像作为输入,因此我们只为宽度和高度指定一个值。scale
和 mean
参数用于将数据缩放到 SSD 训练时使用的值范围。该类的唯一方法是 get_blob
,它接收一个图像并返回一个 blob——一个用于神经网络处理的特殊结构。为了获得 blob,图像首先被调整到指定的正方形大小。然后使用 OpenCV DNN 模块的 blobFromImage
方法,使用指定的 scale
、size
和 mean
值,从调整大小的图像创建 blob。
请注意 get_blob
方法开头的代码。这段代码实现了一个小“技巧”:我们裁剪非方形图像,只提取图像的中心方形部分,如下图所示。
这个技巧是为了保持图像的纵横比不变。如果宽度/高度比发生变化,图像就会变形,目标检测的精度就会降低。这个技巧的一个缺点是,我们只能检测图像中心方形部分(在上图中以蓝色显示)中的人物。
现在让我们来看看使用 SSD 模型进行人员检测的主类。
class SSD:
def __init__(self, frame_proc, ssd_net):
self.proc = frame_proc
self.net = ssd_net
def detect(self, frame):
blob = self.proc.get_blob(frame)
self.net.setInput(blob)
detections = self.net.forward()
# detected object count
k = detections.shape[2]
obj_data = []
for i in np.arange(0, k):
obj = detections[0, 0, i, :]
obj_data.append(obj)
return obj_data
def get_object(self, frame, data):
confidence = int(data[2]*100.0)
(h, w, c) = frame.shape
r_x = int(data[3]*h)
r_y = int(data[4]*h)
r_w = int((data[5]-data[3])*h)
r_h = int((data[6]-data[4])*h)
if w>h :
dx = int((w-h)/2)
r_x = r_x+dx
obj_rect = (r_x, r_y, r_w, r_h)
return (confidence, obj_rect)
def get_objects(self, frame, obj_data, class_num, min_confidence):
objects = []
for (i, data) in enumerate(obj_data):
obj_class = int(data[1])
obj_confidence = data[2]
if obj_class==class_num and obj_confidence>=min_confidence :
obj = self.get_object(frame, data)
objects.append(obj)
return objects
上述类的构造函数有两个参数:用于将图像转换为 blob 的 frame_proc
和用于检测对象的 ssd_net
。主方法 detect
接收一个帧(图像)作为输入,并使用指定的帧处理器从帧中获取一个 blob。blob 用作网络的输入,我们使用 forward
方法获得检测结果。这些检测结果表示为 4 阶数组(张量)。我们不会分析整个张量;我们只需要数组的第二维。我们将它从检测结果中提取出来,并返回结果——一个对象数据列表。
对象数据是一个实际的数组。这是一个例子
[array([ 0. , 15. , 0.90723044, 0.56916684, 0.6017439 ,
0.68543154, 0.93739873], dtype=float32)]
该数组包含七个数字
- 检测到的对象数量
- 其类别的数量
- 对象属于该类别的置信度
- 对象的 ROI 的四条左上角和右下角坐标(坐标相对于 blob 大小)
该类的第二个方法将检测数据转换为更简单的格式以供进一步使用。它将相对置信度转换为百分比值,并将相对 ROI 坐标转换为整数数据——原始图像中的像素坐标。该方法考虑了 blob 数据是从原始帧的中心方形部分提取的事实。
最后,get_objects
方法仅从检测数据中提取具有指定类别且置信度足够的特定对象。因为 DNN 模型可以检测二十种不同类别的对象,所以我们必须仅过滤掉 person
类别,以确保检测到的对象确实是人类,因此我们指定了较高的置信度阈值。
再一个实用类——用于将检测到的对象绘制到图像中以可视化结果。
class Utils:
@staticmethod
def draw_object(obj, label, color, frame):
(confidence, (x1, y1, w, h)) = obj
x2 = x1+w
y2 = y1+h
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
y3 = y1-12
text = label + " " + str(confidence)+"%"
cv2.putText(frame, text, (x1, y3), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1, cv2.LINE_AA)
@staticmethod
def draw_objects(objects, label, color, frame):
for (i, obj) in enumerate(objects):
Utils.draw_object(obj, label, color, frame)
现在我们可以编写代码来启动人员检测算法了。
proto_file = r"C:\PI_RPD\mobilenet.prototxt"
model_file = r"C:\PI_RPD\mobilenet.caffemodel"
ssd_net = CaffeModelLoader.load(proto_file, model_file)
print("Caffe model loaded from: "+model_file)
proc_frame_size = 300
# frame processor for MobileNet
ssd_proc = FrameProcessor(proc_frame_size, 1.0/127.5, 127.5)
person_class = 15
ssd = SSD(ssd_proc, ssd_net)
im_dir = r"C:\PI_RPD\test_images"
im_name = "woman_640x480_01.png"
im_path = os.path.join(im_dir, im_name)
image = cv2.imread(im_path)
print("Image read from: "+im_path)
obj_data = ssd.detect(image)
persons = ssd.get_objects(image, obj_data, person_class, 0.5)
person_count = len(persons)
print("Person count on the image: "+str(person_count))
Utils.draw_objects(persons, "PERSON", (0, 0, 255), image)
res_dir = r"C:\PI_RPD\results"
res_path = os.path.join(res_dir, im_name)
cv2.imwrite(res_path, image)
print("Result written to: "+res_path)
该代码实现了一个 size
=300 的帧处理器,因为我们使用的模型处理大小为 300 x 300 像素的图像。scale
和 mean
参数与用于 MobileNet 模型训练的值相同。这些值必须始终分配给模型的训练值,否则模型的精度会下降。person_class
值为 15,因为在模型上下文中,human 是第 15 个类别。
在样本图像上运行代码会产生以下结果。
我们使用了非常简单的案例来检测人员。目标只是检查我们的代码是否正常工作,以及 DNN 模型是否能够预测图像中是否存在人员。事实证明,它可以工作!
后续步骤
下一步是在 Raspberry Pi 设备上运行我们的代码。在下一篇文章中,我们将介绍如何在设备上安装 Python-OpenCV 并运行代码。