65.9K
CodeProject 正在变化。 阅读更多。
Home

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

2020年10月7日

CPOL

5分钟阅读

viewsIcon

9593

downloadIcon

69

在本文中,我们将展示用于启动这些模型的 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 处理的输入数据的大小。用于图像处理的卷积网络几乎总是使用方形图像作为输入,因此我们只为宽度和高度指定一个值。scalemean 参数用于将数据缩放到 SSD 训练时使用的值范围。该类的唯一方法是 get_blob,它接收一个图像并返回一个 blob——一个用于神经网络处理的特殊结构。为了获得 blob,图像首先被调整到指定的正方形大小。然后使用 OpenCV DNN 模块的 blobFromImage 方法,使用指定的 scalesizemean 值,从调整大小的图像创建 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)]

该数组包含七个数字

  1. 检测到的对象数量
  2. 其类别的数量
  3. 对象属于该类别的置信度
  4. 对象的 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 像素的图像。scalemean 参数与用于 MobileNet 模型训练的值相同。这些值必须始终分配给模型的训练值,否则模型的精度会下降。person_class 值为 15,因为在模型上下文中,human 是第 15 个类别。

在样本图像上运行代码会产生以下结果。

我们使用了非常简单的案例来检测人员。目标只是检查我们的代码是否正常工作,以及 DNN 模型是否能够预测图像中是否存在人员。事实证明,它可以工作!

后续步骤

下一步是在 Raspberry Pi 设备上运行我们的代码。在下一篇文章中,我们将介绍如何在设备上安装 Python-OpenCV 并运行代码。

© . All rights reserved.