使用预训练 SSD 模型检测害虫

• 下载源代码 - 34.8 KB
• 下载模型 - 26 MB

引言

不受控制的野生动物对于企业和房主来说都是一个难题。像鹿、驼鹿甚至猫这样的动物都会对花园、庄稼和财产造成破坏。

在本系列文章中，我们将演示如何使用树莓派实时（或接近实时）地检测害虫（例如驼鹿），然后采取行动驱除害虫。 由于我们不想造成任何伤害，我们将重点关注通过播放响亮的噪音来吓跑害虫。

欢迎您下载项目的源代码。 我们假设您熟悉 Python 并对神经网络的工作原理有基本的了解。

在系列文章的前一篇文章中，我们比较了两种可用于检测害虫的 DNN 类型：检测器和分类器。 检测器获胜。 在本文中，我们将开发使用预训练检测 DNN 检测害虫的 Python 代码。

选择网络架构

对象检测有几种常见的网络架构，例如Faster-RCNN、单次检测器 (SSD) 和You Only Look Once (YOLO)。

由于我们的网络需要在内存和 CPU 有限的边缘设备上运行，我们将使用MobileNet 单次检测器 (SSD) 架构。 MobileNet SSD 是一种轻量级对象检测器网络，在移动和边缘设备上表现良好。 它是在 Pascal VOC 2012 数据集上训练的，该数据集包含一些可能代表害虫的类别，例如猫、牛、狗、马和羊。

我们将使用与之前的系列文章中用于人类检测的算法相同的算法来检测视频中的害虫。

害虫检测代码

首先，我们需要修改 MobileNet 代码以使其检测害虫。

让我们从创建一些实用程序类开始，以使这项任务更容易

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, False, False)
        return blob

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)

CaffeModelLoader 类使用提供的原型和模型文件路径从磁盘加载 Caffe 模型。

下一个实用程序类 FrameProcessor 将帧转换为 blob（用作 CNN 输入的特殊结构化数据）。

最后，Utils 类在帧中检测到的任何对象周围绘制边界矩形。 我们的实用程序类使用的大部分方法都来自OpenCV 库的 Python 版本。 让我们详细看看这些。

这就是我们的实用程序类。 接下来，我们将编写实际检测害虫的代码。

我们将从 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

该类中的关键方法是 detect 和 get_objects。

detect 方法将加载的 DNN 模型应用于每个帧，以检测所有可能类别的对象。

get_objects 方法查看检测到的对象，仅选择属于指定类且被正确检测的概率高（置信度）的对象。

然后，我们将使用 VideoSSD 类，该类对整个视频片段进行害虫检测

class VideoSSD:    
    def __init__(self, ssd):
        self.ssd = ssd
    
    def detect(self, video, class_num, min_confidence, class_name):
        detection_num = 0;
        capture = cv2.VideoCapture(video)
        img = None

        dname = 'Pest detections'
        cv2.namedWindow(dname, cv2.WINDOW_NORMAL)
        cv2.resizeWindow(dname, 1280, 960)
       
        # Capture all frames
        while(True):    
            (ret, frame) = capture.read()
            if frame is None:
                break
        
            obj_data = self.ssd.detect(frame)
            class_objects = self.ssd.get_objects(frame, obj_data, class_num, min_confidence)
            p_count = len(class_objects)
            detection_num += p_count
            
            if len(class_objects)>0:
                Utils.draw_objects(class_objects, class_name, (0, 0, 255), frame)
            
            # Display the resulting frame
            cv2.imshow(dname,frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            
        capture.release()
        cv2.destroyAllWindows()    
        
        return detection_num

该类中唯一的方法是 detect。 它处理从视频文件中提取的所有帧。 在每一帧中，它检测由 class_num 参数指定的类的所有对象，然后在对象周围显示带有边界矩形的帧。

它有效吗？

让我们启动我们的代码，看看它如何处理视频文件。 以下代码加载一个视频文件并尝试检测狗

proto_file = r"C:\PI_PEST\net\mobilenet.prototxt"
model_file = r"C:\PI_PEST\net\mobilenet.caffemodel"
ssd_net = CaffeModelLoader.load(proto_file, model_file)

mobile_proc_frame_size = 300
ssd_proc = FrameProcessor(mobile_proc_frame_size, 1.0/127.5, 127.5)

pest_class = 12
pest_name = "DOG"

ssd = SSD(ssd_proc, ssd_net)

video_file = r"C:\PI_PEST\video\dog_1.mp4"

video_ssd = VideoSSD(ssd)
detections = video_ssd.detect(video_file, pest_class, 0.2, pest_name)

我们将 pest_class 的值设置为 12，因为“狗”是 MobileNet SSD 模型中的第 12 个类别。 这是运行上述代码时捕获的视频。

它会在边缘设备上运行吗？

如您所见，我们的 SSD 检测器在 PC 上运行时成功检测到视频中的狗。 边缘设备呢？ 检测器是否会快速处理数据流以实时检测对象？ 我们可以通过测试帧速率来找出答案，帧速率以每秒帧数 (FPS) 为单位衡量。

在我们之前引用的文章中，我们借用的模型在树莓派 3 设备上的运行速度约为 1.25 FPS。 这足以检测害虫吗？ 我们可以假设，平均而言，动物会在相机上被捕捉至少 2 到 3 秒。 这意味着我们将有 2 到 3 帧来检测害虫并对其做出反应。 听起来几率不错。

后续步骤

到目前为止，结果对于野生动物检测来说并不乐观……但我们不要放弃！

在下一篇文章中，我们将讨论一些检测“异国”害虫（例如驼鹿和犰狳）的想法。